fixing a device proble

matsengrp · Jan 14, 2025 · d484895 · d484895
1 parent fddcee1
commit d484895
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/netam/dcsm.py b/netam/dcsm.py
@@ -347,7 +347,7 @@ def predictions_of_batch(self, batch):
         # However, we have to unsqueeze because scatter_ requires the `index`
         # tensor and the `src` value(s) to have the same shape for broadcasting.
 
-        parent_indices = batch["codon_parents_idxs"]  # Shape: [B, L]
+        parent_indices = batch["codon_parents_idxs"].to(self.device)  # Shape: [B, L]
 
         # Create mask for valid (non-ambiguous) codons.
         valid_mask = parent_indices != AMBIGUOUS_CODON_IDX  # Shape: [B, L]