Update handler.py

handler.py

@@ -210,6 +210,19 @@ class EndpointHandler:
                     frames = self.image_to_video(**generation_kwargs).frames
                 else:
                     frames = self.text_to_video(**generation_kwargs).frames
+                
+                # Log original shape
+                logger.info(f"Original frames shape: {frames.shape}")
+                
+                # Remove batch dimension if present
+                if len(frames.shape) == 5:
+                    frames = frames.squeeze(0)  # Remove batch dimension
+                
+                logger.info(f"Processed frames shape: {frames.shape}")
+                
+                # Ensure we have the correct shape
+                if len(frames.shape) != 4:
+                    raise ValueError(f"Expected tensor of shape [frames, channels, height, width], got shape {frames.shape}")
 
                 # Post-process frames
                 loop = asyncio.new_event_loop()