Rick458 commited on
Commit
5e122e3
1 Parent(s): 27239c5

fixing cuda error

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -105,7 +105,7 @@ to_device(model, device);
105
 
106
  # loading the model
107
  ckp_path = 'indianFood-resnet50.pth'
108
- model.load_state_dict(torch.load(ckp_path), map_location=torch.device('cpu'))
109
  model.eval()
110
 
111
 
 
105
 
106
  # loading the model
107
  ckp_path = 'indianFood-resnet50.pth'
108
+ model.load_state_dict(torch.load(ckp_path), map_location={'cuda': 'cpu'})
109
  model.eval()
110
 
111