update to demo.py

demo.py

@@ -33,6 +33,10 @@ class VideoModel(nn.Module):
         super(VideoModel, self).__init__()
         self.cfg = load_cfg(config)
         self.model = self.build_model()
+        use_gpu = torch.cuda.is_available()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if use_gpu:
+            self.model = self.model.to(self.device)
         self.templates = ['{}']
         self.dataset = self.cfg['data']['dataset']
         self.eval()
@@ -156,7 +160,7 @@ class VideoCLSModel(VideoModel):
             truncation=True,
             max_length=self.model_cfg.max_txt_l.video, 
             return_tensors="pt",
-        )
+        ).to(self.device)
         _, class_embeddings = self.model.encode_text(embeddings)
         return class_embeddings
 
@@ -170,9 +174,7 @@ class VideoCLSModel(VideoModel):
             images = values[0]
             target = values[1]
 
-            if torch.cuda.is_available():
-                images = images.cuda(non_blocking=True)
-                target = target.cuda(non_blocking=True)
+            images = images.to(self.device)
 
             # encode images
             images = rearrange(images, 'b c k h w ->  b k c h w')
@@ -190,7 +192,7 @@ class VideoCLSModel(VideoModel):
             similarity = self.model.get_sim(image_features, self.text_features)[0]
 
             all_outputs.append(similarity.cpu())
-            all_targets.append(target.cpu())
+            all_targets.append(target)
 
         all_outputs = torch.cat(all_outputs)
         all_targets = torch.cat(all_targets)