fix tests by moving to float32

misko · Jun 30, 2024 · c15745f · c15745f
1 parent 07c2857
commit c15745f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/spf/dataset/spf_dataset.py b/spf/dataset/spf_dataset.py
@@ -449,9 +449,9 @@ def render_session(self, receiver_idx, session_idx):
             pd = torch_get_phase_diff(data["signal_matrix"]).to(torch.float32)
             data["x"] = torch.vstack([abs_signal[0], abs_signal[1], pd])[None]
 
-        data["y_rad"] = data["ground_truth_theta"][None]
-        data["y_phi"] = data["ground_truth_phi"][None]
-        data["craft_y_rad"] = data["craft_ground_truth_theta"][None]
+        data["y_rad"] = data["ground_truth_theta"][None].to(torch.float32)
+        data["y_phi"] = data["ground_truth_phi"][None].to(torch.float32)
+        data["craft_y_rad"] = data["craft_ground_truth_theta"][None].to(torch.float32)
 
         # data["y_discrete"] = v5_thetas_to_targets(data["y_rad"], self.nthetas)