tests and precommit

microsoft · Jul 30, 2024 · b3405f0 · b3405f0
1 parent bfb86ff
commit b3405f0
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
@@ -69,6 +69,7 @@ def move_to_cpu(tensor_list):
     for tensor in tensor_list:
         tensor.data = tensor.data.cpu()
 
+
 @contextmanager
 def unwrap_model_for_generation(model):
     """

diff --git a/tests/unit/runtime/zero/test_unwrap_model.py b/tests/unit/runtime/zero/test_unwrap_model.py
@@ -38,11 +38,13 @@
 elif get_accelerator().is_bf16_supported():
     config["bf16"] = {"enabled": True}
 
+
 class TestUnwrapModel(DistributedTest):
     # gather across more than 1 gpu
     world_size = 2
 
     def test(self):
+
         def hooks_exist(engine):
             if engine.optimizer is not None and hasattr(engine.optimizer, "parameter_offload"):
                 optimizer_offload = engine.optimizer.parameter_offload