Update src/accelerate/accelerator.py

Co-authored-by: Zach Mueller <[email protected]>
huggingface · Nov 7, 2023 · 5419ad2 · 5419ad2
1 parent eda082f
commit 5419ad2
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/src/accelerate/accelerator.py b/src/accelerate/accelerator.py
@@ -2580,9 +2580,8 @@ def save_model(
 
         os.makedirs(save_directory, exist_ok=True)
 
-        for param in model.parameters():
-            if param.device == torch.device("meta"):
-                raise RuntimeError("You can't save the model since some parameters are on the meta device.")
+        if any(param.device == torch.device("meta") for param in model.parameters()):
+            raise RuntimeError("You can't save the model since some parameters are on the meta device.")
 
         # get the state_dict of the model
         state_dict = self.get_state_dict(model)