Bring scores back to cpu.

mila-iqia · Dec 26, 2024 · d0d0b23 · d0d0b23
1 parent 9da231e
commit d0d0b23
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/diffusion_for_multi_scale_molecular_dynamics/analysis/score_viewer.py b/src/diffusion_for_multi_scale_molecular_dynamics/analysis/score_viewer.py
@@ -153,7 +153,7 @@ def _compute_projected_scores(self, score_network: ScoreNetwork):
         for time, sigma in zip(self.times, self.sigmas):
             batch = self._get_batch(time, sigma)
 
-            sigma_normalized_scores = score_network(batch).X.detach()
+            sigma_normalized_scores = score_network(batch).X.detach().cpu()
             vectors = einops.rearrange(
                 sigma_normalized_scores, "batch natoms space -> batch (natoms space)"
             )