fix: get rid of reduce mocking (for testing)

tests/test_gym.py

@@ -1,10 +1,9 @@
-from unittest.mock import call, patch
+from unittest.mock import call
 
 import torch
 
 from modalities.batch import DatasetBatch
 from modalities.gym import Gym
-from modalities.running_env.fsdp.reducer import Reducer
 
 
 def test_run_cpu_only(
@@ -37,15 +36,13 @@ def test_run_cpu_only(
     llm_data_loader_mock.__len__ = lambda _: num_batches
 
     gym = Gym(trainer=trainer, evaluator=evaluator_mock, loss_fun=loss_mock, num_ranks=num_ranks)
-    with patch.object(Reducer, "reduce", return_value=None) as reduce_mock:
-        gym.run(
-            model=nn_model_mock,
-            optimizer=optimizer_mock,
-            callback_interval_in_batches=int(num_batches),
-            train_data_loader=llm_data_loader_mock,
-            evaluation_data_loaders=[],
-            checkpointing=checkpointing_mock,
-        )
-        nn_model_mock.forward.assert_has_calls([call(b.samples) for b in batches])
-        optimizer_mock.step.assert_called()
-        reduce_mock.assert_called()
+    gym.run(
+        model=nn_model_mock,
+        optimizer=optimizer_mock,
+        callback_interval_in_batches=int(num_batches),
+        train_data_loader=llm_data_loader_mock,
+        evaluation_data_loaders=[],
+        checkpointing=checkpointing_mock,
+    )
+    nn_model_mock.forward.assert_has_calls([call(b.samples) for b in batches])
+    optimizer_mock.step.assert_called()