From 15b5da2e1a2bd9d339efc6b428ac4274d186309f Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Sat, 10 Dec 2022 04:51:45 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 5f178fd2..9ac2b164 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -435,11 +435,17 @@ class DaViTStage(nn.Module): self.blocks = SequentialWithSize(*stage_blocks) def forward(self, x : Tensor, size: Tuple[int, int]): - x : Tensor, size: Tuple[int, int] = self.patch_embed(x, size) + x, size = self.patch_embed(x, size) + x : Tensor = x + size : Tuple[int, int] = size if self.grad_checkpointing and not torch.jit.is_scripting(): - x : Tensor, size: Tuple[int, int] = checkpoint_seq(self.blocks, x, size) + x, size = checkpoint_seq(self.blocks, x, size) else: - x : Tensor, size: Tuple[int, int] = self.blocks(x, size) + x, size = self.blocks(x, size) + + x : Tensor = x + size : Tuple[int, int] = size + return x, size