bug fix for batch size other than 1 in offset diversity.

microsoft · Feb 28, 2024 · 14fa240 · 14fa240
1 parent 16e5085
commit 14fa240
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/DCVC-DC/src/models/video_model.py b/DCVC-DC/src/models/video_model.py
@@ -53,8 +53,8 @@ def forward(self, x, aux_feature, flow):
         # warp
         offset = offset.view(B * self.group_num * self.offset_num, 2, H, W)
         mask = mask.view(B * self.group_num * self.offset_num, 1, H, W)
-        x = x.view(B * self.group_num, C // self.group_num, H, W)
-        x = x.repeat(self.offset_num, 1, 1, 1)
+        x = x.repeat(1, self.offset_num, 1, 1)
+        x = x.view(B * self.group_num * self.offset_num, C // self.group_num, H, W)
         x = flow_warp(x, offset)
         x = x * mask
         x = x.view(B, C * self.offset_num, H, W)