Update hawk/models/video_llama.py

hawk/models/video_llama.py

@@ -403,11 +403,11 @@ class VideoLLAMA(Blip2Base):
 
     def encode_videoQformer_visual(self, image, motion=False):
         if motion is False:
-            device = image.device
+            device = "cuda:0"
             
             # input shape b,c,t,h,w
             batch_size,_,time_length,_,_ = image.size()
-            image = einops.rearrange(image, 'b c t h w -> (b t) c h w')
+            image = einops.rearrange(image, 'b c t h w -> (b t) c h w').to(device)
             with self.maybe_autocast():
                 
                 # embed image features with blip2, out: (b t) q h
@@ -454,12 +454,12 @@ class VideoLLAMA(Blip2Base):
         
         else:
             # Motion Encoder
-            device = image.device
+            device = "cuda:0"
             
             # input shape b,c,t,h,w
             batch_size,_,time_length,_,_ = image.size()
             
-            image = einops.rearrange(image, 'b c t h w -> (b t) c h w')
+            image = einops.rearrange(image, 'b c t h w -> (b t) c h w').to(device)
             
             with self.maybe_autocast():