Make sure we are using CUDA if available

app.py

@@ -13,16 +13,14 @@ from torchvision.transforms import (
     Lambda,
     Resize, RandomCrop,
 )
-from transformers import pipeline
-
+from transformers import VideoMAEForVideoClassification, VideoMAEFeatureExtractor
 
 MODEL_CKPT = "omermazig/videomae-finetuned-nba-5-class-4-batch-8000-vid-multiclass"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 CLIPS_FROM_SINGLE_VIDEO = 5
 
-pipe = pipeline("video-classification", model=MODEL_CKPT)
-trained_model = pipe.model
-image_processor = pipe.image_processor
+trained_model = VideoMAEForVideoClassification.from_pretrained(MODEL_CKPT).to(DEVICE)
+image_processor = VideoMAEFeatureExtractor.from_pretrained(MODEL_CKPT)
 
 mean = image_processor.image_mean
 std = image_processor.image_std
@@ -67,7 +65,7 @@ def parse_video_to_clips(video_file):
         video_clips_list.append(inference_transform(video_clip))
 
     videos_tensor = torch.stack([single_clip.permute(1, 0, 2, 3) for single_clip in video_clips_list])
-    return videos_tensor
+    return videos_tensor.to(DEVICE)
 
 
 def infer(video_file):