diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py
index 4e5c23b8ef..41a9222abe 100644
--- a/tests/post_training/unit/train_distill_test.py
+++ b/tests/post_training/unit/train_distill_test.py
@@ -1047,7 +1047,6 @@ def __call__(self, input_tokens, **kwargs):
     if hasattr(trainer2.checkpoint_manager, "wait_until_finished"):
       trainer2.checkpoint_manager.wait_until_finished()
 
-  @mock.patch("maxtext.trainers.post_train.distillation.train_distill.distillation_utils.OfflineArrayRecordIterator")
   @mock.patch("maxtext.trainers.post_train.distillation.train_distill.MaxTextDistillationTrainer")
   @mock.patch("maxtext.trainers.post_train.distillation.train_distill.input_pipeline_interface.create_data_iterator")
   @mock.patch("maxtext.trainers.post_train.distillation.train_distill.get_maxtext_model")
@@ -1062,7 +1061,6 @@ def test_main_offline_mode_skips_teacher_loading(
       mock_get_model,
       mock_create_iterator,
       mock_trainer_cls,
-      mock_offline_iter_cls,
   ):
     """Verifies offline mode (offline_data_dir is set) skips teacher model loading."""
     # 1. Configs