test fix

Author: kazeevn

src/wyckoff_transformer/tests/test_adaln_conditioning.py

@@ -152,7 +152,9 @@ def test_forward_with_cond_returns_correct_shape(self):
         model = self._build_model(condition_dim=1)
         start = torch.tensor([0, 1, 2], dtype=torch.int64)
         cascade = [torch.tensor([[0, 1], [1, 0], [2, 0]], dtype=torch.int64)]
-        cond = torch.tensor([[0.5], [-0.5], [1.5]], dtype=torch.float32)
+        # Match the model's parameter dtype: torch.set_default_dtype may have been
+        # changed by another test (e.g. MACE calculator init flips it to float64).
+        cond = torch.tensor([[0.5], [-0.5], [1.5]], dtype=next(model.parameters()).dtype)
         out = model(start, cascade, padding_mask=None, prediction_head=0, cond=cond)
         # Output is the prediction head applied to the last token; shape [batch_size, num_classes]
         self.assertEqual(out.shape, (3, 3))