Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent b799ba95e2
commit 01bcbc94f7

@ -36,7 +36,7 @@ __all__ = ['DaViT']
class SequentialWithSize(nn.Sequential):
def forward(self, x : Tensor, size: Tuple[int, int]):
for module in self._modules.values():
for module in self:
x, size = module(x, size)
return x, size

Loading…
Cancel
Save