Update app.py

app.py

@@ -12,10 +12,22 @@ os.environ['USE_TF'] = '0'  # Set TensorFlow to off
 os.environ['USE_TORCH'] = '1'  # Set PyTorch to on
 
 from doctr.models import ocr_predictor
-ocr_model = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Example for PyTorch model or doctr model
 
+ocr_model = ocr_predictor(det_arch='db_resnet50', reco_arch='crnn_vgg16_bn', pretrained=True).to(device)
 
+
+import torch
+
+# Check if CUDA is available
+if torch.cuda.is_available():
+    print(f"GPU is available. Device: {torch.cuda.get_device_name(0)}")
+else:
+    print("GPU is not available, using CPU.")
+
+    
 """
 Perform OCR with doctr
 """