Make sure random_tokens is on correct device (#96)

Make sure random_tokens is on correct device
lucidrains · Aug 23, 2021 · 7c516af · 7c516af
1 parent c85682e
commit 7c516af
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/alphafold2_pytorch/mlm.py b/alphafold2_pytorch/mlm.py
@@ -70,7 +70,7 @@ def noise(self, seq, mask):
 
         random_replace_token_prob_mask = get_mask_subset_with_prob(mlm_mask, (1 - self.keep_token_same_prob) * self.random_replace_token_prob)
 
-        random_tokens = torch.randint(1, constants.NUM_AMINO_ACIDS, seq.shape)
+        random_tokens = torch.randint(1, constants.NUM_AMINO_ACIDS, seq.shape).to(seq.device)
 
         for token_id in self.exclude_token_ids:
             random_replace_token_prob_mask = random_replace_token_prob_mask & (random_tokens != token_id)  # make sure you never substitute a token with an excluded token type (pad, start, end)
@@ -89,4 +89,4 @@ def forward(self, seq_embed, original_seq, mask):
         seq_labels = original_seq[mask]
 
         loss = F.cross_entropy(seq_logits, seq_labels, reduction = 'mean')
-        return loss
+        return loss