logo
Browse Source

Update device

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

2
nn_fingerprint.py

@ -69,7 +69,7 @@ class Model:
def __call__(self, *args, **kwargs): def __call__(self, *args, **kwargs):
new_args = [] new_args = []
new_kwargs = {} new_kwargs = {}
for x in new_args:
for x in args:
new_args.append(x.to(self.device)) new_args.append(x.to(self.device))
for k, v in kwargs.items(): for k, v in kwargs.items():
new_kwargs[k] = v.to(self.device) new_kwargs[k] = v.to(self.device)

Loading…
Cancel
Save