Upload modeling_fastesm.py with huggingface_hub

modeling_fastesm.py

@@ -717,11 +717,11 @@ class EmbeddingMixin:
                     seqs = to_embed[i * batch_size:(i + 1) * batch_size]
                     input_ids, attention_mask = batch['input_ids'].to(device), batch['attention_mask'].to(device)
                     residue_embeddings = self._embed(input_ids, attention_mask)
-                    embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype).cpu()
+                    embeddings = get_embeddings(residue_embeddings, attention_mask).to(embed_dtype)
                     for seq, emb, mask in zip(seqs, embeddings, attention_mask):
                         if full_embeddings:
                             emb = emb[mask.bool()].reshape(-1, hidden_size)
-                        embeddings_dict[seq] = emb
+                        embeddings_dict[seq] = emb.cpu()
 
         if save:
             torch.save(embeddings_dict, save_path)