fixing broken device swap

crystal_diffusion/models/score_network.py

@@ -175,8 +175,7 @@ def _forward_unchecked(self, batch: Dict[AnyStr, torch.Tensor]) -> torch.Tensor:
             computed_scores : the scores computed by the model.
         """
         positions = batch[self.position_key]  # shape [batch_size, number_of_atoms, spatial_dimension]
-        times = batch[self.timestep_key]  # shape [batch_size, 1]
-
+        times = batch[self.timestep_key].to(positions.device)  # shape [batch_size, 1]
         input = torch.cat([self.flatten(positions), times], dim=1)
 
         output = self.mlp_layers(input).reshape(positions.shape)

crystal_diffusion/samplers/noisy_position_sampler.py

@@ -79,6 +79,6 @@ def get_noisy_position_sample(real_relative_positions: torch.Tensor, sigmas: tor
             "sigmas array is expected to be of the same shape as the real_relative_positions array"
 
         z_scores = NoisyPositionSampler._get_gaussian_noise(real_relative_positions.shape)
-        noise = sigmas * z_scores
+        noise = (sigmas * z_scores).to(real_relative_positions.device)
         noisy_relative_positions = map_positions_to_unit_cell(real_relative_positions + noise)
         return noisy_relative_positions

crystal_diffusion/score/wrapped_gaussian_score.py

@@ -32,8 +32,8 @@
 import numpy as np
 import torch
 
-SIGMA_THRESHOLD = 1.0 / np.sqrt(2.0 * np.pi)
-U_THRESHOLD = 0.5
+SIGMA_THRESHOLD = torch.Tensor([1.0 / np.sqrt(2.0 * np.pi)])
+U_THRESHOLD = torch.Tensor([0.5])
 
 
 def get_sigma_normalized_score_brute_force(u: float, sigma: float, kmax: Optional[int] = None) -> float:
@@ -124,7 +124,8 @@ def get_sigma_normalized_score(
     for mask_calculator, score_calculator in zip(mask_calculators, score_calculators):
         mask = mask_calculator(list_u, list_sigma)
         if mask.any():
-            flat_view[mask] = score_calculator(list_u[mask], list_sigma[mask], list_k)
+            device = flat_view.device
+            flat_view[mask] = score_calculator(list_u[mask], list_sigma.to(device)[mask], list_k.to(device))
 
     return sigma_normalized_scores
 
@@ -139,7 +140,8 @@ def _get_small_sigma_small_u_mask(list_u: torch.Tensor, list_sigma: torch.Tensor
     Returns:
         mask_1a : an array of booleans of shape [Nu]
     """
-    return torch.logical_and(list_sigma <= SIGMA_THRESHOLD, list_u < U_THRESHOLD)
+    device = list_u.device
+    return torch.logical_and(list_sigma.to(device) <= SIGMA_THRESHOLD.to(device), list_u < U_THRESHOLD.to(device))
 
 
 def _get_small_sigma_large_u_mask(list_u: torch.Tensor, list_sigma: torch.Tensor) -> torch.Tensor:
@@ -152,7 +154,8 @@ def _get_small_sigma_large_u_mask(list_u: torch.Tensor, list_sigma: torch.Tensor
     Returns:
         mask_1b : an array of booleans of shape [Nu]
     """
-    return torch.logical_and(list_sigma <= SIGMA_THRESHOLD, list_u >= U_THRESHOLD)
+    device = list_u.device
+    return torch.logical_and(list_sigma.to(device) <= SIGMA_THRESHOLD.to(device), list_u >= U_THRESHOLD.to(device))
 
 
 def _get_large_sigma_mask(list_u: torch.Tensor, list_sigma: torch.Tensor) -> torch.Tensor: