update return msg

model.py

@@ -78,6 +78,7 @@ class HybridCNNViT(nn.Module):
 
 
 def load_and_pad_single_image(image_path, img_size=(224, 224)):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     img = cv2.imread(image_path)
     if img is None:
         raise ValueError(f"Could not read image: {image_path}")
@@ -145,7 +146,7 @@ def check_file(image_path):
         return checkpoint
 
     model = HybridCNNViT(3, 2)
-    checkpoint = torch.load('checkpoint32.pth')
+    checkpoint = torch.load("/home/user/app/checkpoint32.pth")
     checkpoint = remove_module_from_checkpoint(checkpoint)
     model.load_state_dict(checkpoint['model_state_dict'])
     # optimizer.load_state_dict(checkpoint['optimizer_state_dict'])