ctheodoris commited on
Commit
57f02a4
1 Parent(s): 25dd1da

clone embs_i to resolve memory leak in cls embs

Browse files
Files changed (1) hide show
  1. geneformer/emb_extractor.py +1 -1
geneformer/emb_extractor.py CHANGED
@@ -148,7 +148,7 @@ def get_embs(
148
  del embs_h
149
  del dict_h
150
  elif emb_mode == "cls":
151
- cls_embs = embs_i[:,0,:] # CLS token layer
152
  embs_list.append(cls_embs)
153
  del cls_embs
154
 
 
148
  del embs_h
149
  del dict_h
150
  elif emb_mode == "cls":
151
+ cls_embs = embs_i[:,0,:].clone().detach() # CLS token layer
152
  embs_list.append(cls_embs)
153
  del cls_embs
154