Update rex_chest_xray .py

12341123 · Apr 23, 2021 · dfac4ed · dfac4ed
1 parent 07dbec0
commit dfac4ed
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/rex_chest_xray .py b/rex_chest_xray .py
@@ -528,6 +528,7 @@ def inference(name, model, device, data_loader, criterion, limit=None):
 
 model = densenet121(num_classes=len(datasets.default_pathologies))
 model.load_state_dict(torch.load("rex_chest_xray/nih_mc_cx-densenet121-valid-best.pt").state_dict())
+model = model.to(device)
 
 test_auc, test_loss, task_aucs = inference(name='Test',
                                          model=model,