Update handler.py

handler.py

@@ -2,9 +2,12 @@ from typing import  Dict, List, Any
 from transformers import AutoTokenizer
 from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig
 import torch
+from loguru import logger
 
 # check for GPU
-#device = 0 if torch.cuda.is_available() else -1
+device = 0 if torch.cuda.is_available() else -1
+
+logger.info(f"cuda: {device}")
 
 MAX_INPUT_TOKEN_LENGTH  = 4000
 MAX_MAX_NEW_TOKENS      = 2048
@@ -32,8 +35,7 @@ class EndpointHandler():
         if input_token_length > MAX_INPUT_TOKEN_LENGTH:
             return [{"generated_text": None, "error": f"input is too long ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH})"}] 
 
-        #input_ids = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
-        input_ids = self.tokenizer(inputs, return_tensors="pt").input_ids
+        input_ids = self.tokenizer(inputs, return_tensors="pt").to(self.model.device)
 
         outputs = self.model.generate(**input_ids, **parameters)