Update modeling_vila.py

modeling_vila.py

@@ -739,15 +739,18 @@ class VILAForCausalLM(VILAPretrainedModel):
                     self.encoders[name].pool_sizes[0][0] = 4 * round_up_to_bucket(num_video_frames / 256)
 
                 if num_video_frames > 512:
-                    media_split = []
-                    frames_split = 4
+                    media_split, num_splits = [], []
                     for video in media[name]:
-                        media_split += video.tensor_split(frames_split, dim=0)
+                        video_split = video.split(512, dim=0)
+                        media_split.extend(video_split)
+                        num_splits.append(len(video_split))
+
                     embeds_split = []
                     for video in media_split:
                         embeds_split += self.encoders[name]([video], media_config[name])
+
                     embeds_merged = [
-                        torch.cat(embeds_split[i * frames_split: (i + 1) * frames_split], dim=0)
+                        torch.cat(embeds_split[i * num_splits[i]: (i + 1) * num_splits[i]], dim=0)
                         for i in range(len(media[name]))
                     ]
                     embeds[name] = deque(embeds_merged)