testing with saved best checkpoint

2025-03-17 22:20:44 +00:00
parent 4ba69f2d1b
commit 91915d6a14
1 changed files with 1 additions and 1 deletions
@@ -370,7 +370,7 @@ def main(args, criterion):

        if epoch == (args.epochs - 1):
            checkpoint = torch.load(os.path.join(args.output_dir, args.task, 'checkpoint-best.pth'), map_location='cpu')
-            model.load_state_dict(checkpoint['model'], strict=False)
+            model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
            model.to(device)
            print("Test with the best model, epoch = %d:" % checkpoint['epoch'])
            test_stats, auc_roc = evaluate(data_loader_test, model, device, args, -1, mode='test',