From 19f405255c09169257d2ead69389de3c0f413e8e Mon Sep 17 00:00:00 2001
From: Christian Medina <37550954+cmedinasoriano@users.noreply.github.com>
Date: Wed, 1 Jul 2026 16:50:48 -0400
Subject: [PATCH] remove auto device map in training strategies

---
 train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index 1406b83..6991b6c 100644
--- a/train.py
+++ b/train.py
@@ -41,7 +41,7 @@ def train(config_path):
         model = AutoModelForCausalLM.from_pretrained(
             config["base_model"],
             torch_dtype=torch.float16,
-            device_map="auto",
+            # device_map="auto",
             trust_remote_code=True,
         )
         print("✓ Success: 4-bit AS-IS")
@@ -81,7 +81,7 @@ def train(config_path):
                     model = AutoModelForCausalLM.from_pretrained(
                         config["base_model"],
                         torch_dtype=torch.bfloat16,
-                        device_map="auto",
+                        # device_map="auto",
                         low_cpu_mem_usage=True,
                         trust_remote_code=True,
                     )