Enforce max_memory for device_map strategies (#22311)

Enforce for device_map strategies
huggingface · Mar 23, 2023 · 62d9baa · 62d9baa
1 parent 6828768
commit 62d9baa
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -2576,6 +2576,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
                     model,
                     dtype=torch_dtype,
                     low_zero=(device_map == "balanced_low_0"),
+                    max_memory=max_memory,
                     **kwargs,
                 )
             kwargs["max_memory"] = max_memory