bert feats to device

Graylab · Mar 3, 2023 · ead1f5a · ead1f5a
1 parent 6483845
commit ead1f5a
Showing 1 changed file with 1 addition and 8 deletions.
diff --git a/igfold/model/IgFold.py b/igfold/model/IgFold.py
@@ -220,14 +220,7 @@ def forward(
 
         ### Model forward pass
 
-        # bert_feats, bert_attns = [], []
-        # for t in tokens:
-        #     f, a, h = self.get_bert_feats(t)
-        #     bert_feats.append(f)
-        #     bert_attns.append(a)
-        #     bert_hidden.append(h)
-
-        bert_feats = torch.cat(embeddings, dim=1)
+        bert_feats = torch.cat(embeddings, dim=1).to(self.device)
         bert_attn = torch.zeros(
             (batch_size, seq_len, seq_len, self.bert_attn_dim),
             device=self.device,