From b2dfa605577163ab6b447541bb26cc1bed317e26 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Thu, 8 Dec 2022 03:01:25 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 2bcc2b9b..3ed13849 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -410,12 +410,12 @@ class DaViT(nn.Module): for i in range(self.num_stages)]) #main_blocks = [] - self.main_blocks = nn.ModuleList() - for block_id, block_param in enumerate(self.architecture): - layer_offset_id = len(list(itertools.chain(*self.architecture[:block_id]))) + self.stages = nn.Sequential() + for stage_id, stage_param in enumerate(self.architecture): + layer_offset_id = len(list(itertools.chain(*self.architecture[:stage_id]))) - block = nn.ModuleList([ - nn.ModuleList([ + stage = nn.Sequential([ + nn.Sequential([ ChannelBlock( dim=self.embed_dims[item], num_heads=self.num_heads[item], @@ -438,12 +438,12 @@ class DaViT(nn.Module): window_size=window_size, ) if attention_type == 'spatial' else None for attention_id, attention_type in enumerate(attention_types)] - ) for layer_id, item in enumerate(block_param) + ) for layer_id, item in enumerate(stage_param) ]) - self.main_blocks.add_module(f'block', block) + self.main_blocks.append(stage) - self.feature_info += [dict(num_ch=self.embed_dims[block_id], reduction = 2, module=f'block.{block_id}')] + self.feature_info += [dict(num_ch=self.embed_dims[stage_id], reduction = 2, module=f'stage.{stage_id}')] #self.main_blocks = nn.ModuleList(main_blocks) '''