"git@developer.sourcefind.cn:chenpangpang/transformers.git" did not exist on "f16ff0f07e3867db5feda00a661572a190f404e6"
Unverified Commit 5fd4e3c8 authored by Sylvain Gugger's avatar Sylvain Gugger Committed by GitHub
Browse files

Enforce `max_memory` for device_map strategies (#22311)

Enforce  for device_map strategies
parent 48bef3a7
...@@ -2571,6 +2571,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix ...@@ -2571,6 +2571,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMix
model, model,
dtype=torch_dtype, dtype=torch_dtype,
low_zero=(device_map == "balanced_low_0"), low_zero=(device_map == "balanced_low_0"),
max_memory=max_memory,
**kwargs, **kwargs,
) )
kwargs["max_memory"] = max_memory kwargs["max_memory"] = max_memory
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment