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],