enable GPU use

app.py

@@ -5,11 +5,13 @@ import plotly.express as px
 import torch
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 model_checkpoint = "facebook/xlm-v-base"
 
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
 model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
+model = model.to(device)
 mask_token = tokenizer.mask_token
 
 
@@ -27,7 +29,7 @@ def eval_prob(target_word, text):
     target_idx = tokenizer.encode(target_word)[-2]
 
     # Convert masked text to token IDs
-    inputs = tokenizer(text_masked, return_tensors="pt")
+    inputs = tokenizer(text_masked, return_tensors="pt").to(device)
     
     # Calculate logits score (for each token, for each position)
     token_logits = model(**inputs).logits