Enable cuda

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
3 years ago · 6023656152
1 changed files with 2 additions and 2 deletions
--- a/pytorchvideo.py
+++ b/pytorchvideo.py
@ -65,8 +65,8 @@ class PytorchVideo(NNOperator):
                self.classmap[v] = str(k).replace('"', '')
        else:
            self.classmap = classmap
-        # todo: self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
-        self.device = 'cpu'
+        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.eval()