From b322b563eb2eeddd5db40a4a02d34cba258d36bd Mon Sep 17 00:00:00 2001 From: liuht Date: Wed, 27 Dec 2023 01:22:21 +0800 Subject: [PATCH] [FIX] Fix GPU Memory Leak Issue in hetero_spec Generation --- .../market/heterogeneous/organizer/hetero_map/__init__.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/learnware/market/heterogeneous/organizer/hetero_map/__init__.py b/learnware/market/heterogeneous/organizer/hetero_map/__init__.py index b2f39fe..8968793 100644 --- a/learnware/market/heterogeneous/organizer/hetero_map/__init__.py +++ b/learnware/market/heterogeneous/organizer/hetero_map/__init__.py @@ -287,6 +287,9 @@ class HeteroMap(nn.Module): # go through transformers, get the first cls embedding encoder_output = self.encoder(**outputs) # bs, seqlen+1, hidden_dim output_features = encoder_output[:, 0, :] + + del inputs, outputs, encoder_output + torch.cuda.empty_cache() return output_features @@ -316,6 +319,8 @@ class HeteroMap(nn.Module): with torch.no_grad(): output_features = self._extract_features(bs_x_test).detach().cpu().numpy() output_feas_list.append(output_features) + del output_features + torch.cuda.empty_cache() all_output_features = np.concatenate(output_feas_list, 0) return all_output_features