From b799ba95e2bc5920a2a2227f352e6c173da9a5b7 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Sat, 10 Dec 2022 00:20:05 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 51162c5a..0b8c219f 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -486,7 +486,7 @@ class DaViT(nn.Module): stage_drop_rates = dpr[len(attention_types) * sum(depths[:stage_id]):len(attention_types) * sum(depths[:stage_id + 1])] stage = DaViTStage( - in_chans if stage_id == 0 else embed_dims[i - 1], + in_chans if stage_id == 0 else embed_dims[stage_id - 1], embed_dims[stage_id], depth = 1, patch_size = patch_size,