Params4bit added to bnb classes in set_module_tensor_to_device() (#…

…2315)
huggingface · Jan 10, 2024 · 456afd9 · 456afd9
1 parent 0d2280d
commit 456afd9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/accelerate/utils/modeling.py b/src/accelerate/utils/modeling.py
@@ -328,7 +328,7 @@ def set_module_tensor_to_device(
             param is not None
             and param.device.type != "cuda"
             and torch.device(device).type == "cuda"
-            and param_cls.__name__ in ["Int8Params", "FP4Params"]
+            and param_cls.__name__ in ["Int8Params", "FP4Params", "Params4bit"]
         ):
             device_quantization = device
             device = "cpu"