Update app.py

app.py

@@ -5,7 +5,7 @@ from PIL import Image
 from torchvision import transforms
 import gradio as gr
 import numpy as np
-
+import spaces
 
 from pytorch_grad_cam.utils.image import show_cam_on_image
 from scripts.trainer import XrayReg
@@ -93,6 +93,7 @@ def predict_custom(model, input_tensor):
     return float(pred)
 
 
+@spaces.GPU
 def predict_and_cam_custom(inp, model):
     input_tensor, rgb_img = preprocess_image_custom(inp)
     value = predict_custom(model, input_tensor)
@@ -169,8 +170,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
-
-
-if __name__ == "__main__":
-    main()
+    main()