From 1488fe198b351d99ec629dc4cc455dd53943a03d Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Thu, 28 Jul 2022 11:52:39 +0800 Subject: [PATCH] Change checkpoint parent folder Signed-off-by: Jael Gu --- nn_fingerprint.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/nn_fingerprint.py b/nn_fingerprint.py index fd6e651..c891256 100644 --- a/nn_fingerprint.py +++ b/nn_fingerprint.py @@ -70,7 +70,7 @@ class NNFingerprint(NNOperator): log.info('Loading weights...') if checkpoint_path is None: path = str(Path(__file__).parent) - checkpoint_path = os.path.join(path, './checkpoints/pfann_fma_m.pt') + checkpoint_path = os.path.join(path, 'saved_model', 'pfann_fma_m.pt') state_dict = torch.load(checkpoint_path, map_location=self.device) self.model.load_state_dict(state_dict) self.model.eval()