From 60236561527e13ad2091e65f62cbf42969d4f1ef Mon Sep 17 00:00:00 2001
From: Jael Gu <mengjia.gu@zilliz.com>
Date: Thu, 16 Jun 2022 11:14:53 +0800
Subject: [PATCH] Enable cuda

Signed-off-by: Jael Gu <mengjia.gu@zilliz.com>
---
 pytorchvideo.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pytorchvideo.py b/pytorchvideo.py
index 5808e77..7aa32bd 100644
--- 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()