device

matsengrp · Jan 15, 2025 · 4d2debf · 4d2debf
1 parent d52a617
commit 4d2debf
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/netam/dcsm.py b/netam/dcsm.py
@@ -272,12 +272,13 @@ def predictions_of_batch(self, batch):
         # the full batch. Namely, we need to set the parent codon to 1 -
         # sum(children).
 
+        # TODO I guess we could store indicator in self and have everything move with a self.to(device) call.
         # This indicator lifts things up from aa land to codon land.
         indicator = CODON_AA_INDICATOR_MATRIX.to(self.device).T
         log_preds = (
             log_neutral_codon_probs
             + log_selection_factors @ indicator
-            + self.stop_codon_zapper
+            + self.stop_codon_zapper.to(self.device)
         )
         assert torch.isnan(log_preds).sum() == 0