From 6b3a2669887015ca7db074b04d0a6773410712bc Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Tue, 6 Dec 2022 22:12:05 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index ae5694d4..af61eda0 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -507,7 +507,8 @@ class DaViT(nn.Module): branches = [0] for block_index, block_param in enumerate(self.architecture): - branch_ids = sorted(set(block_param)) + branch_ids = sorted(set(torch.Tensor(block_param))) + #branch_ids = [block_index] for branch_id in branch_ids: if branch_id not in branches: x, size = self.patch_embeds[branch_id](features[-1], sizes[-1])