Update davit.py

3 years ago · c82f9204e7
parent a813bbd1df
commit c82f9204e7
1 changed files with 2 additions and 1 deletions
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@ -514,7 +514,8 @@ class DaViT(nn.Module):
            features[-1], sizes[-1] = patch_layer(features[-1], sizes[-1])
            
            for block in enumerate(stage):
-                for _, layer in enumerate(block):
+                for layer in enumerate(block):
+                    print(layer)
                    if self.grad_checkpointing and not torch.jit.is_scripting():
                        features[-1], sizes[-1] = checkpoint.checkpoint(layer, (features[-1], sizes[-1]))
                    else: