From 2308ee4bdf818180167dc853264ebefb36f85a2b Mon Sep 17 00:00:00 2001
From: Fredo Guan <fredo.guan@hotmail.com>
Date: Sat, 10 Dec 2022 04:49:46 -0800
Subject: [PATCH] Update davit.py

---
 timm/models/davit.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/timm/models/davit.py b/timm/models/davit.py
index d74f7361..5f178fd2 100644
--- a/timm/models/davit.py
+++ b/timm/models/davit.py
@@ -435,11 +435,11 @@ class DaViTStage(nn.Module):
         self.blocks = SequentialWithSize(*stage_blocks)
         
     def forward(self, x : Tensor, size: Tuple[int, int]):
-        x, size = self.patch_embed(x, size)
+        x : Tensor, size: Tuple[int, int] = self.patch_embed(x, size)
         if self.grad_checkpointing and not torch.jit.is_scripting():
-            x, size = checkpoint_seq(self.blocks, x, size)
+            x : Tensor, size: Tuple[int, int] = checkpoint_seq(self.blocks, x, size)
         else:
-            x, size = self.blocks(x, size)
+            x : Tensor, size: Tuple[int, int] = self.blocks(x, size)
         return x, size