Update app

app.py

@@ -85,6 +85,21 @@ def load_image(image_file, input_size=448, max_num=12):
     pixel_values = torch.stack(pixel_values)
     return pixel_values
 
+# Save the original cuda() method
+original_cuda = torch.Tensor.cuda
+
+# Define a new cuda() method
+def safe_cuda(self, *args, **kwargs):
+    if torch.cuda.is_available():
+        return original_cuda(self, *args, **kwargs)  # Use the original cuda() method
+    else:
+        return self  # Return the tensor itself (stays on CPU)
+
+# Monkey-patch the cuda() method
+torch.Tensor.cuda = safe_cuda
+
+
+
 
 model_name = "YuukiAsuna/Vintern-1B-v2-ViTable-docvqa"