From 87abcafaa573b5c85a41a815b4104ef79b88775c Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Thu, 8 Dec 2022 05:54:17 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index f8f61d73..16cc99f4 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -38,7 +38,7 @@ class MySequential(nn.Sequential): def forward(self, inputs : Tensor, size : Tuple[int, int]): for module in self: - inputs = module(inputs, size) + inputs, size = module(inputs, size) return inputs class ConvPosEnc(nn.Module): @@ -185,7 +185,7 @@ class ChannelBlock(nn.Module): act_layer=act_layer) - def forward(self, x, size: Tuple[int, int]): + def forward(self, x : Tensor, size: Tuple[int, int]): x = self.cpe[0](x, size) cur = self.norm1(x) cur = self.attn(cur) @@ -310,7 +310,7 @@ class SpatialBlock(nn.Module): act_layer=act_layer) - def forward(self, x, size: Tuple[int, int]): + def forward(self, x : Tensor, size: Tuple[int, int]): H, W = size B, L, C = x.shape