MingGatsby commited on
Commit
0117147
1 Parent(s): bf8ba85

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -148,7 +148,7 @@ def load_model(root_dir, model_name, model_file_name):
148
  model = Build_Custom_Model(model_name, NUM_CLASSES, pretrained=False).to(device)
149
  else:
150
  model = SEResNet50(spatial_dims=2, in_channels=1, num_classes=NUM_CLASSES).to(device)
151
- model.load_state_dict(torch.load(os.path.join(root_dir, model_file_name), map_location=device)))
152
  model.eval()
153
  return model
154
 
 
148
  model = Build_Custom_Model(model_name, NUM_CLASSES, pretrained=False).to(device)
149
  else:
150
  model = SEResNet50(spatial_dims=2, in_channels=1, num_classes=NUM_CLASSES).to(device)
151
+ model.load_state_dict(torch.load(os.path.join(root_dir, model_file_name), map_location=device))
152
  model.eval()
153
  return model
154