logo
Browse Source

Fix device issue when save_model

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

2
nn_fingerprint.py

@ -157,7 +157,7 @@ class NNFingerprint(NNOperator):
path = os.path.join(path, name)
dummy_input = torch.rand(
(1,) + (self.params['n_mels'], self.params['u'])
)
).to(self.device)
if format == 'pytorch':
path = path + '.pt'
torch.save(self.model, path)

Loading…
Cancel
Save