From b9af370cbe9084e58ad5e10529c7f7fd44739df4 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Sat, 10 Dec 2022 04:59:53 -0800 Subject: [PATCH] Update davit.py --- timm/models/davit.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/timm/models/davit.py b/timm/models/davit.py index 9ac2b164..a2a315c6 100644 --- a/timm/models/davit.py +++ b/timm/models/davit.py @@ -438,10 +438,12 @@ class DaViTStage(nn.Module): x, size = self.patch_embed(x, size) x : Tensor = x size : Tuple[int, int] = size + dim1 : int = size[0] + dim2 : int = size[1] if self.grad_checkpointing and not torch.jit.is_scripting(): - x, size = checkpoint_seq(self.blocks, x, size) + x, size = checkpoint_seq(self.blocks, x, (dim1, dim2)) else: - x, size = self.blocks(x, size) + x, size = self.blocks(x, (dim1, dim2)) x : Tensor = x size : Tuple[int, int] = size @@ -592,7 +594,7 @@ class DaViT(nn.Module): - def forward_features(self, x): + def forward_features(self, x : Tensor): #x, sizes = self.forward_network(x) size: Tuple[int, int] = (x.size(2), x.size(3)) x, size = self.stages(x, size)