MingGatsby
commited on
Commit
•
0117147
1
Parent(s):
bf8ba85
Update app.py
Browse files
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 |
|