Merge pull request #4091 from AI-Hypercomputer:bvandermoon-test-fixes

PiperOrigin-RevId: 927552045

Author: Google-ML-Automation

tests/post_training/unit/train_distill_test.py

@@ -1064,6 +1064,8 @@ def test_main_offline_mode_skips_teacher_loading(
   ):
     """Verifies offline mode (offline_data_dir is set) skips teacher model loading."""
     # 1. Configs
+    devices = jax.devices()[:1]
+    mock_create_mesh.return_value = np.array(devices)
     mock_global = mock.Mock()
     mock_global.student_overrides = {}
     mock_global.teacher_overrides = {}  # No checkpoint needed
@@ -1172,6 +1174,8 @@ def test_main_online_mode_loads_teacher(
       mock_trainer_cls,
   ):
     """Verifies online mode (offline_data_dir is None) loads both student and teacher models."""
+    devices = jax.devices()[:1]
+    mock_create_mesh.return_value = np.array(devices)
     mock_global = mock.Mock()
     mock_global.student_overrides = {}
     mock_global.teacher_overrides = {"load_parameters_path": "gs://ckpt"}