Browse Source
Fix device issue when save_model
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
1 additions and
1 deletions
-
nn_fingerprint.py
|
@ -157,7 +157,7 @@ class NNFingerprint(NNOperator): |
|
|
path = os.path.join(path, name) |
|
|
path = os.path.join(path, name) |
|
|
dummy_input = torch.rand( |
|
|
dummy_input = torch.rand( |
|
|
(1,) + (self.params['n_mels'], self.params['u']) |
|
|
(1,) + (self.params['n_mels'], self.params['u']) |
|
|
) |
|
|
|
|
|
|
|
|
).to(self.device) |
|
|
if format == 'pytorch': |
|
|
if format == 'pytorch': |
|
|
path = path + '.pt' |
|
|
path = path + '.pt' |
|
|
torch.save(self.model, path) |
|
|
torch.save(self.model, path) |
|
|