From feb3f51765d1d4af5f21e0cb16f33fddc1061bb1 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Thu, 8 Dec 2022 03:03:48 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 3ed13849..6e165894 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -414,8 +414,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 = nn.Sequential( + nn.Sequential( ChannelBlock( dim=self.embed_dims[item], num_heads=self.num_heads[item], @@ -437,9 +437,9 @@ class DaViT(nn.Module): cpe_act=cpe_act, window_size=window_size, ) if attention_type == 'spatial' else None - for attention_id, attention_type in enumerate(attention_types)] + for attention_id, attention_type in enumerate(attention_types) ) for layer_id, item in enumerate(stage_param) - ]) + ) self.main_blocks.append(stage)