Update modeling_fastesm.py

modeling_fastesm.py

@@ -447,8 +447,8 @@ class FastEsmPreTrainedModel(PreTrainedModel):
                 for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), desc='Embedding batches'):
                     seqs = sequences[i * batch_size:(i + 1) * batch_size]
                     input_ids, attention_mask = batch['input_ids'].to(device), batch['attention_mask'].to(device)
-                    residue_embeddings = self.forward(input_ids, attention_mask, output_hidden_states=True).hidden_states[-1].float() # required for sql
-                    embeddings = get_embeddings(residue_embeddings, attention_mask)
+                    residue_embeddings = self.forward(input_ids, attention_mask, output_hidden_states=True).hidden_states[-1].detach().float() # required for sql
+                    embeddings = get_embeddings(residue_embeddings, attention_mask).cpu()
 
                     for seq, emb in zip(seqs, embeddings):
                         c.execute("INSERT OR REPLACE INTO embeddings VALUES (?, ?)", 
@@ -466,10 +466,10 @@ class FastEsmPreTrainedModel(PreTrainedModel):
             for i, batch in tqdm(enumerate(dataloader), total=len(dataloader), desc='Embedding batches'):
                 seqs = sequences[i * batch_size:(i + 1) * batch_size]
                 input_ids, attention_mask = batch['input_ids'].to(device), batch['attention_mask'].to(device)
-                residue_embeddings = self.forward(input_ids, attention_mask, output_hidden_states=True).hidden_states[-1].float()
+                residue_embeddings = self.forward(input_ids, attention_mask, output_hidden_states=True).hidden_states[-1].detach().float()
                 if full_precision:
                     residue_embeddings = residue_embeddings.float()
-                embeddings = get_embeddings(residue_embeddings, attention_mask)
+                embeddings = get_embeddings(residue_embeddings, attention_mask).cpu()
                 for seq, emb in zip(seqs, embeddings):
                     embeddings_dict[seq] = emb