Change the device of the tokenizer's output

functions.py

@@ -29,6 +29,7 @@ def summarize_text(text: str):
     print(['summarize_text', 'start'])
     input_text = f'<s>Instruction: Elabora un resume del siguiente texto.\nInput: {text}\nOutput: '
     batch = tokenizer(input_text, return_tensors='pt')
+    batch = batch.to('cuda')
     print(['summarize_text', 'generating'])
     with torch.cuda.amp.autocast():
         output_tokens = model.generate(**batch,
@@ -45,8 +46,10 @@ def load_model(peft_model_id):
     print(['load_model', 'start'])
     config = PeftConfig.from_pretrained(peft_model_id)
     print(['load_model', 'loading model'])
-    model = AutoModelForCausalLM.from_pretrained(
-        config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='balanced')
+    model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path,
+                                                 return_dict=True,
+                                                 load_in_8bit=True,
+                                                 device_map='auto')
     print(['load_model', 'loading tokenizer'])
     tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
     model = PeftModel.from_pretrained(model, peft_model_id)