Upload app.py with huggingface_hub

app.py

@@ -13,8 +13,9 @@ model_path = hf_hub_download(
 
 # ✅ Load the Model
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = models.resnet50(pretrained=False)
+print(f"✅ Using device: {device}")
 
+model = models.resnet50(pretrained=False)
 checkpoint = torch.load(model_path, map_location=device)
 
 if "model" in checkpoint:
@@ -24,6 +25,7 @@ elif "state_dict" in checkpoint:
 else:
     model.load_state_dict(checkpoint, strict=False)
 
+# Ensure correct output layer (21 classes for Clothing1M)
 model.fc = torch.nn.Linear(2048, 21)
 model.to(device)
 model.eval()
@@ -42,7 +44,8 @@ def preprocess_image(image):
         transforms.ToTensor(),
         transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
     ])
-    return transform(image).unsqueeze(0).to(device)
+    image = transform(image).unsqueeze(0).to(device)  # Ensure tensor is on GPU if available
+    return image
 
 # ✅ Prediction Function
 def predict(image):