Update handler.py

handler.py

@@ -7,6 +7,10 @@ class EndpointHandler:
         # Configuração do modelo
         self.model_name_or_path = "souzat19/Llama3.1_fn14133.29122024"
         
+        # Detecta se GPU está disponível
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        print(f"Using device: {self.device}")
+        
         print("Initializing tokenizer...")
         # Inicialização do tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(
@@ -18,10 +22,9 @@ class EndpointHandler:
         # Inicialização do modelo com configurações mínimas
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name_or_path,
-            device_map=None,  # Desativa mapeamento automático
             torch_dtype=torch.float32,
             trust_remote_code=True
-        ).cpu()  # Força o uso de CPU
+        ).to(self.device)  # Move para GPU se disponível
         
         print("Model initialized successfully")
         
@@ -51,6 +54,9 @@ Você é um assistente especializado em planejamento de compras públicas de aco
                 truncation=True,
                 max_length=4096
             )
+            
+            # Move input para mesmo device do modelo
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
             # Gera a resposta
             with torch.no_grad():