From b3526cc9a4b0ba692ba43b0dfa9ecde2ba1cee28 Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Tue, 6 Dec 2022 22:35:45 -0800
Subject: [PATCH] Update davit.py

---
 timm/models/davit.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index fb7c2ff2..da494f68 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -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: