Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent c1a6a2fe17
commit bb4d5f74f7

@ -37,7 +37,7 @@ __all__ = ['DaViT']
# modified nn.Sequential that includes a size tuple in the forward function # modified nn.Sequential that includes a size tuple in the forward function
class SequentialWithSize(nn.Sequential): class SequentialWithSize(nn.Sequential):
@torch.jit.ignore @torch.jit.ignore
def forward(self, x : Tensor, size: Tuple[int, int]): def forward(self, x : Tensor, size: Tuple[int, int]):
for module in self._modules.values(): for module in self._modules.values():
x, size = module(x, size) x, size = module(x, size)

Loading…
Cancel
Save