From 258de970aa71c13568e92ff6b70fade699fcc587 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Thu, 16 Jun 2022 11:04:51 +0800 Subject: [PATCH] Disable gpu since fail for some model Signed-off-by: Jael Gu --- pytorchvideo.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pytorchvideo.py b/pytorchvideo.py index 86c32bd..5808e77 100644 --- a/pytorchvideo.py +++ b/pytorchvideo.py @@ -65,7 +65,8 @@ class PytorchVideo(NNOperator): self.classmap[v] = str(k).replace('"', '') else: 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.eval()