made changes made based on checkpoint model

app.py

@@ -30,13 +30,13 @@ device = torch.device('cpu')
 model = models.resnet50(pretrained=False)  # Load the ResNet-50 architecture
 model.fc = nn.Linear(model.fc.in_features, 1000) 
 
-model.load_state_dict(torch.load("model.pth", map_location=device))  # Load the trained weights (.pth)
+# model.load_state_dict(torch.load("model.pth", map_location=device))  # Load the trained weights (.pth)
 model.to(device)  # Move model to CPU (even if you have a GPU)
 
 checkpoint = torch.load('model.pth', map_location='cpu')
 
 # Load the model weights
-model.load_state_dict(checkpoint['model_state_dict'], strict=False)
+model.load_state_dict(checkpoint['model_state_dict'], strict=False, map_location=device)
 
 # If you need to resume training, load optimizer and scheduler states
 optimizer.load_state_dict(checkpoint['optimizer_state_dict'])