Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent 9d9959919b
commit 721487c304

@ -35,10 +35,10 @@ from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
__all__ = ['DaViT'] __all__ = ['DaViT']
class SequentialWithSize(nn.Sequential): class SequentialWithSize(nn.Sequential):
def forward(self, x : Tensor, size: Tuple[int, int]): def forward(self, x: Tensor, size: Tuple[int, int]):
for module in self: for module in self:
output = module(inputs, size) output = module(x, size)
inputs : Tensor = output[0] x : Tensor = output[0]
size : Tuple[int, int] = output[1] size : Tuple[int, int] = output[1]
return x, size return x, size

Loading…
Cancel
Save