Fix batched inference

interpolator.py

@@ -125,8 +125,8 @@ def debug_forward(self, x0, x1, batch_dt) -> Dict[str, List[torch.Tensor]]:
         # Note: In film_net we fix time to be 0.5, and recursively invoke the interpo-
         # lator for multi-frame interpolation. Below, we create a constant tensor of
         # shape [B]. We use the `time` tensor to infer the batch size.
-        backward_flow = util.multiply_pyramid(backward_flow_pyramid, batch_dt[:, 0])
-        forward_flow = util.multiply_pyramid(forward_flow_pyramid, 1 - batch_dt[:, 0])
+        backward_flow = util.multiply_pyramid(backward_flow_pyramid, batch_dt)
+        forward_flow = util.multiply_pyramid(forward_flow_pyramid, 1 - batch_dt)
 
         pyramids_to_warp = [
             util.concatenate_pyramids(image_pyramids[0][:self.fusion_pyramid_levels],
@@ -154,6 +154,5 @@ def debug_forward(self, x0, x1, batch_dt) -> Dict[str, List[torch.Tensor]]:
             'backward_flow_pyramid': backward_flow_pyramid,
         }
 
-    @torch.jit.export
     def forward(self, x0, x1, batch_dt) -> torch.Tensor:
         return self.debug_forward(x0, x1, batch_dt)['image'][0]

util.py

@@ -102,7 +102,7 @@ def multiply_pyramid(pyramid: List[torch.Tensor],
     # the batch of images from BxHxWxC-format to CxHxWxB. This can then be
     # multiplied with a batch of scalars, then we transpose back to the standard
     # BxHxWxC form.
-    return [image * scalar for image in pyramid]
+    return [image * scalar[..., None, None] for image in pyramid]
 
 
 def flow_pyramid_synthesis(