Update main.py

main.py

@@ -19,12 +19,9 @@ feature_extractor = ViTFeatureExtractor.from_pretrained(model_name)
 # Load the trained model weights
 num_classes = 7
 model.classifier = nn.Linear(model.config.hidden_size, num_classes)
-model.load_state_dict(torch.load("/kaggle/input/skincancer-vit/skin_cancer_model.pth", map_location=torch.device('cpu')))
+model.load_state_dict(torch.load("skin_cancer_model.pth", map_location=torch.device('cpu')))
 model.eval()
 
-# Define device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
 
 # Define class labels
 class_labels = ['benign_keratosis-like_lesions', 'basal_cell_carcinoma', 'actinic_keratoses', 'vascular_lesions', 'melanocytic_Nevi', 'melanoma', 'dermatofibroma']
@@ -43,7 +40,7 @@ transform = transforms.Compose([
 async def predict(file: UploadFile = File(...)):
     contents = await file.read()
     image = Image.open(io.BytesIO(contents))
-    image = transform(image).unsqueeze(0).to(device)  # Add batch dimension and move to device
+    image = transform(image).unsqueeze(0)  # Add batch dimension and move to device
     
     with torch.no_grad():
         outputs = model(image)