Fix: Use GPU

inference.py

@@ -8,6 +8,7 @@ import torch
 from collections import OrderedDict
 
 # Third Party Imports
+import spaces
 from torchvision import transforms
 from torch.nn import functional as F
 from torchvision.models import resnet50
@@ -16,6 +17,7 @@ from pytorch_grad_cam.utils.image import show_cam_on_image
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 
 
+@spaces.GPU
 def inference(image, alpha, top_k, target_layer, model=None, classes=None):
     """
     Function to run inference on the input image