From 600a23fa23e25d9628840695c8ba55f01111ebff Mon Sep 17 00:00:00 2001 From: Gabriele Corso Date: Tue, 18 Jul 2023 01:40:55 -0700 Subject: [PATCH] Fix ckpt path #85 --- confidence/dataset.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/confidence/dataset.py b/confidence/dataset.py index 1103b79..f9f660d 100644 --- a/confidence/dataset.py +++ b/confidence/dataset.py @@ -213,7 +213,7 @@ class ConfidenceDataset(Dataset): t_to_sigma = partial(t_to_sigma_compl, args=self.original_model_args) model = get_model(self.original_model_args, self.device, t_to_sigma=t_to_sigma, no_parallel=True) - state_dict = torch.load(f'{self.original_model_dir}/{self.model_ckpt}.pt', map_location=torch.device('cpu')) + state_dict = torch.load(f'{self.original_model_dir}/{self.model_ckpt}', map_location=torch.device('cpu')) model.load_state_dict(state_dict, strict=True) model = model.to(self.device) model.eval()