Update src/backend/model_loader.py

src/backend/model_loader.py

@@ -73,20 +73,21 @@ def load_model(
         model = AutoModelForCausalLM.from_pretrained( model_name, **args)
         
         # Move modelo para CPU se não há GPU e device_map não foi usado
-        if device_map is None and not has_cuda:
-            model = model.to("cpu")
+        #if device_map is None and not has_cuda:
+        #    model = model.to("cpu")
         
         # Cria pipeline
         pipeline_kwargs = {
             "model": model,
             "tokenizer": tokenizer,
+            'device_map': args["device_map"]
         }
         
         # Só adiciona device ao pipeline se não usar device_map no modelo
-        if device_map is None:
-            pipeline_kwargs["device"] = 0 if has_cuda else -1
-        else:
-            pipeline_kwargs["device_map"] = device_map
+        #if device_map is None:
+        #    pipeline_kwargs["device"] = 0 if has_cuda else -1
+        #else:
+        #    pipeline_kwargs["device_map"] = device_map
         
         pipe = pipeline("text-generation", **pipeline_kwargs)