From e7f94cbf18c9bb7402968fa2583ec90cf3a58fb0 Mon Sep 17 00:00:00 2001 From: Alberto Dall'olio Date: Tue, 29 Aug 2023 16:38:27 +0000 Subject: [PATCH] add mask size for more modularity --- core/update.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/update.py b/core/update.py index f940497f..7b83c207 100644 --- a/core/update.py +++ b/core/update.py @@ -112,7 +112,7 @@ def forward(self, net, inp, corr, flow): return net, None, delta_flow class BasicUpdateBlock(nn.Module): - def __init__(self, args, hidden_dim=128, input_dim=128): + def __init__(self, args, hidden_dim=128, input_dim=128, mask_size=8): super(BasicUpdateBlock, self).__init__() self.args = args self.encoder = BasicMotionEncoder(args) @@ -122,7 +122,7 @@ def __init__(self, args, hidden_dim=128, input_dim=128): self.mask = nn.Sequential( nn.Conv2d(128, 256, 3, padding=1), nn.ReLU(inplace=True), - nn.Conv2d(256, 64*9, 1, padding=0)) + nn.Conv2d(256, mask_size**2 * 9, 1, padding=0)) def forward(self, net, inp, corr, flow, upsample=True): motion_features = self.encoder(flow, corr)