logo
Browse Source

Change checkpoint parent folder

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 2 years ago
parent
commit
1488fe198b
  1. 2
      nn_fingerprint.py

2
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()

Loading…
Cancel
Save