Browse Source
Disable gpu since fail for some model
Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
main
1 changed files with
2 additions and
1 deletions
-
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() |
|
|
|
