app

app.py

@@ -19,6 +19,7 @@ def mlm(image, text):
         }
         ans = list()
         res = defaultdict(list)
+        device = 'cpu'
         for k, v in questions_dict.items():
             predicted_tokens = []
             tokenized_text = tokenizer.tokenize(v)
@@ -27,8 +28,9 @@ def mlm(image, text):
             segments_ids = [0] * len(tokenized_text)
             
             # Convert inputs to PyTorch tensors
-            tokens_tensor = torch.tensor([indexed_tokens]).to('cuda')
-            segments_tensors = torch.tensor([segments_ids]).to('cuda')
+        
+            tokens_tensor = torch.tensor([indexed_tokens]).to(device)
+            segments_tensors = torch.tensor([segments_ids]).to(device)
 
             masked_index = tokenized_text.index('[MASK]')
             with torch.no_grad():