Move tensor to device in greedy translation function

app.py

@@ -8,8 +8,9 @@ def en_translate_ar_beam(text, model, tokenizer, max_tries=50):
     return "future work"
 
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 def en_translate_ar_greedy(text, model, tokenizer, max_tries=50):
-    source_tensor = torch.tensor(tokenizer(text)).unsqueeze(0)
+    source_tensor = torch.tensor(tokenizer(text)).unsqueeze(0).to(device)
     target_tokens = greedy_decode(model, source_tensor,
                                   tokenizer.get_tokenId('<s>'),
                                   tokenizer.get_tokenId('</s>'),
@@ -18,7 +19,6 @@ def en_translate_ar_greedy(text, model, tokenizer, max_tries=50):
     return tokenizer.decode(target_tokens)
 
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
 tokenizer = Callable_tokenizer('./assets/tokenizers/en-ar_tokenizer.model')
 
 model_state_dict = torch.load("./assets/models/en-ar_s2sAttention.pth", map_location=device, weights_only=True)['model_state_dict']