logo
Browse Source

Disable gpu since fail for some model

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
Jael Gu 2 years ago
parent
commit
258de970aa
  1. 3
      pytorchvideo.py

3
pytorchvideo.py

@ -65,7 +65,8 @@ class PytorchVideo(NNOperator):
self.classmap[v] = str(k).replace('"', '') self.classmap[v] = str(k).replace('"', '')
else: else:
self.classmap = classmap self.classmap = classmap
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
# todo: self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
self.device = 'cpu'
self.model = torch.hub.load('facebookresearch/pytorchvideo', model=model_name, pretrained=True) self.model = torch.hub.load('facebookresearch/pytorchvideo', model=model_name, pretrained=True)
self.model.eval() self.model.eval()

Loading…
Cancel
Save