From 14fa2408315e76ab67eecc86771a81a2ccc36ca3 Mon Sep 17 00:00:00 2001 From: Bin Li Date: Tue, 27 Feb 2024 14:02:39 +0800 Subject: [PATCH] bug fix for batch size other than 1 in offset diversity. --- DCVC-DC/src/models/video_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/DCVC-DC/src/models/video_model.py b/DCVC-DC/src/models/video_model.py index 7cf5bc7..31388fb 100644 --- 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)