fix: remove deletion of removed model_kwargs key

axolotl-ai-cloud · Dec 26, 2023 · cd34680 · cd34680
1 parent fb442c9
commit cd34680
Showing 1 changed file with 0 additions and 1 deletion.
diff --git a/src/axolotl/utils/models.py b/src/axolotl/utils/models.py
@@ -447,7 +447,6 @@ def load_model(
             model_kwargs["device"] = torch.cuda.current_device()
             del model_kwargs["torch_dtype"]
             del model_kwargs["device_map"]
-            del model_kwargs["max_memory"]
 
             model = MambaLMHeadModel.from_pretrained(
                 base_model,