From 63769a454078a0f05d3aec3c5a6f86cf625e6c50 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Mon, 15 Aug 2022 18:49:46 +0800 Subject: [PATCH] Fix device issue when save_model 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 3623a7d..cef8c08 100644 --- a/nn_fingerprint.py +++ b/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)