Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent 792564543c
commit 6b3a266988

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

Loading…
Cancel
Save