Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent deb3e09b3d
commit b3526cc9a4

@ -506,9 +506,7 @@ class DaViT(nn.Module):
features = [x]
sizes = [size]
branches = [0]
block_index : int
for block_index, block_param in enumerate(self.architecture):
'''
branch_ids = sorted(set(block_param))
@ -520,7 +518,7 @@ class DaViT(nn.Module):
branches.append(branch_id)
'''
block_index : int = block_index
if block_index not in branches:
x, size = self.patch_embeds[block_index](features[-1], sizes[-1])
@ -528,11 +526,10 @@ class DaViT(nn.Module):
sizes.append(size)
branches.append(branch_id)
layer_index : int
branch_id : int
for layer_index, branch_id in enumerate(block_param):
layer_index : int = layer_index
branch_id : int = branch_id
if self.grad_checkpointing and not torch.jit.is_scripting():
features[branch_id], _ = checkpoint.checkpoint(self.main_blocks[block_index][layer_index], features[branch_id], sizes[branch_id])
else:

Loading…
Cancel
Save