From 44928197f6fc1ca735236ac68c86113f83da27ce Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Thu, 8 Dec 2022 05:34:47 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index b7085aec..acb0900a 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -33,6 +33,13 @@ from .registry import register_model __all__ = ['DaViT'] + +class MySequential(nn.Sequential): + def forward(self, inputs : Tensor, size : Tuple[int, int]]): + for module in self._modules.values(): + inputs = module(inputs, size) + return inputs + class ConvPosEnc(nn.Module): def __init__(self, dim : int, k : int=3, act : bool=False, normtype : str='none'): @@ -411,8 +418,8 @@ class DaViT(nn.Module): for stage_id, stage_param in enumerate(self.architecture): layer_offset_id = len(list(itertools.chain(*self.architecture[:stage_id]))) - stage = nn.Sequential(*[ - nn.Sequential(*[ + stage = MySequential(*[ + Sequential(*[ ChannelBlock( dim=self.embed_dims[item], num_heads=self.num_heads[item],