space updated

app.py

@@ -7,7 +7,7 @@ TOKENIZER = T5Tokenizer.from_pretrained('werent4/mt5TranslatorLT')
 MODEL = MT5ForConditionalGeneration.from_pretrained("werent4/mt5TranslatorLT")
 MODEL.to(DEVICE)
 
-def translate(text, device, max_length, num_beams, translation_way = "en-lt"):
+def translate(text, max_length, num_beams, translation_way = "en-lt"):
     translations_ways = {
         "en-lt": "<EN2LT>",
         "lt-en": "<LT2EN>"
@@ -15,7 +15,7 @@ def translate(text, device, max_length, num_beams, translation_way = "en-lt"):
     if translation_way not in translations_ways:
         raise ValueError(f"Invalid translation way. Supported ways: {list(translations_ways.keys())}")
     input_text = f"{translations_ways[translation_way]} {text}"
-    encoded_input = TOKENIZER(input_text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(device)
+    encoded_input = TOKENIZER(input_text, return_tensors="pt", padding=True, truncation=True, max_length=128).to(DEVICE)
     with torch.no_grad():
         output_tokens = MODEL.generate(
           **encoded_input,