MingGatsby
commited on
Commit
•
5063525
1
Parent(s):
a56d0aa
Update app.py
Browse files
app.py
CHANGED
@@ -148,8 +148,6 @@ 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 |
-
print(os.path.join(root_dir, model_file_name))
|
152 |
-
print("=================================")
|
153 |
model.load_state_dict(torch.load(os.path.join(root_dir, model_file_name), map_location=device))
|
154 |
model.eval()
|
155 |
return model
|
|
|
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
|