Fixed device mismatch in energy_spectrum computation

MagriLab · Sep 4, 2023 · 2338b25 · 2338b25
1 parent 21d7c7b
commit 2338b25
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kolsol/torch/solver.py b/kolsol/torch/solver.py
@@ -405,7 +405,7 @@ def energy_spectrum(self, u_hat: torch.Tensor, agg: bool = False) -> torch.Tenso
         uu = 0.5 * oe.contract('...u -> ...', u_hat * torch.conj(u_hat)).real
         intk = torch.sqrt(self.kk).to(torch.int)
 
-        ek = torch.zeros((tuple([*leading_dims]) + tuple([torch.max(intk) + 1])))
+        ek = torch.zeros((tuple([*leading_dims]) + tuple([torch.max(intk) + 1])), device=self.device)
         for combo in it.product(*map(range, intk.shape)):
             ek[tuple([...]) + tuple([intk[combo]])] += uu[tuple([...]) + tuple(combo)] / self.nk_grid ** self.ndim