Spaces:

anjikum
/

resent50_imagenet1k_trained_v1

Sleeping

anjikum commited on Dec 31, 2024

Commit

44d2b34

verified ·

1 Parent(s): 1b084ff

made changes to model loading

Files changed (1) hide show

app.py CHANGED Viewed

@@ -36,7 +36,7 @@ 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'])
 # If you need to resume training, load optimizer and scheduler states
 optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

 checkpoint = torch.load('model.pth', map_location='cpu')
 # Load the model weights
+model.load_state_dict(checkpoint['model_state_dict'], strict=False)
 # If you need to resume training, load optimizer and scheduler states
 optimizer.load_state_dict(checkpoint['optimizer_state_dict'])