Browse Source
Change checkpoint parent folder
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
1 additions and
1 deletions
-
nn_fingerprint.py
|
|
@ -70,7 +70,7 @@ class NNFingerprint(NNOperator): |
|
|
log.info('Loading weights...') |
|
|
log.info('Loading weights...') |
|
|
if checkpoint_path is None: |
|
|
if checkpoint_path is None: |
|
|
path = str(Path(__file__).parent) |
|
|
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) |
|
|
state_dict = torch.load(checkpoint_path, map_location=self.device) |
|
|
self.model.load_state_dict(state_dict) |
|
|
self.model.load_state_dict(state_dict) |
|
|
self.model.eval() |
|
|
self.model.eval() |
|
|
|
