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