Update test_deepspeed.py (#2323)

huggingface · Jan 9, 2024 · 2a8829d · 2a8829d
1 parent 3969731
commit 2a8829d
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/deepspeed/test_deepspeed.py b/tests/deepspeed/test_deepspeed.py
@@ -31,6 +31,7 @@
 
 import accelerate
 from accelerate.accelerator import Accelerator
+from accelerate.scheduler import AcceleratedScheduler
 from accelerate.state import AcceleratorState
 from accelerate.test_utils.testing import (
     AccelerateTestCase,
@@ -363,7 +364,7 @@ def test_prepare_deepspeed(self, optim_type, scheduler_type):
                 self.assertTrue(accelerator.deepspeed_config["train_batch_size"], 16)
                 self.assertEqual(type(model), DeepSpeedEngine)
                 self.assertEqual(type(optimizer), DeepSpeedOptimizerWrapper)
-                self.assertEqual(type(lr_scheduler), DeepSpeedSchedulerWrapper)
+                self.assertEqual(type(lr_scheduler), AcceleratedScheduler)
                 self.assertEqual(type(accelerator.deepspeed_engine_wrapped), DeepSpeedEngineWrapper)
 
         elif optim_type == DS_OPTIMIZER and scheduler_type == DS_SCHEDULER: