fix log

huggingface · Dec 4, 2023 · 223a946 · 223a946
1 parent 32b3bf6
commit 223a946
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/big_modeling.py b/src/accelerate/big_modeling.py
@@ -412,7 +412,7 @@ def dispatch_model(
 
         # warn if there is any params on the meta device
         offloaded_devices_str = " and ".join(
-            [device for device in set(device_map.values()) if device in ("cpu", "disk")]
+            [device for device in set(device_map.values()) if device in offloaded_devices]
         )
         if len(offloaded_devices_str) > 0:
             logging.warning(