Update app.py

app.py

@@ -247,17 +247,14 @@ class Translators:
         model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
             dtype=torch.bfloat16,
-            device_map="auto",
-            max_memory={0: "16GB"},
-        )
+            device_map="auto").to(self.device)
         tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         context = {"tone": "Formal and precise"}        
         system = f"Translate the user's text to {self.tl}. Provide the final translation in a formal tone immediately immediately without any other text."
         messages = [
             {"role": "system", "content": system},
             {"role": "user", "content": self.input_text},
-        ]
-        
+        ]     
         prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         print(prompt)
         inputs = tokenizer(prompt, return_tensors="pt").to(self.device)
@@ -266,8 +263,7 @@ class Translators:
             outputs = model.generate(
                 **inputs,
                 max_new_tokens=self.max_new_tokens,
-            )
-        
+            )   
         generated_tokens = outputs[0][input_length:]
         translation = tokenizer.decode(generated_tokens, skip_special_tokens=True)
         return translation