test: Fix tests

Author: rrutmann

src/modalities/trainer.py

@@ -161,9 +161,7 @@ def _train_batch(
         """
         if scheduled_pipeline is not None:
             if isinstance(loss_fun, ChunkedLMHeadCrossEntropyLoss):
-                raise NotImplementedError(
-                    "ChunkedLMHeadCrossEntropyLoss is not supported with pipeline parallelism."
-                )
+                raise NotImplementedError("ChunkedLMHeadCrossEntropyLoss is not supported with pipeline parallelism.")
             pp_schedule = scheduled_pipeline.pp_schedule
             # Pipeline Parallel forward / backward inside step() call
             # with self.train_context(optional_context_parallel_ctx):
@@ -250,10 +248,10 @@ def train(
             m.train()
 
         local_num_seen_samples = 0
-        cumulated_losses = torch.zeros(3).cuda()
 
         # throughput
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        cumulated_losses = torch.zeros(3, device=device)
 
         # batch loop
         batch: DatasetBatch

tests/conftest.py

@@ -197,7 +197,7 @@ def gradient_clipper_mock() -> GradientClipperIF:
 
 @pytest.fixture(scope="function")
 def loss_mock() -> Loss:
-    return MagicMock(spec=Loss, return_value=torch.rand(1, requires_grad=True))
+    return MagicMock(spec=Loss, return_value=torch.rand((), requires_grad=True))
 
 
 @pytest.fixture(scope="function")

tests/test_weight_tying.py

@@ -97,7 +97,7 @@ def _make_norm_config() -> LayerNormWrapperConfig:
 
 def create_device_mesh_stub(*mesh_dim_names: str) -> DeviceMesh:
     device_mesh = DeviceMesh.__new__(DeviceMesh)
-    device_mesh.mesh_dim_names = mesh_dim_names
+    device_mesh._mesh_dim_names = mesh_dim_names
     return device_mesh