Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent 143a239e93
commit 128d9729c3

@ -36,7 +36,7 @@ from .registry import register_model
__all__ = ['DaViT']
class MySequential(nn.Sequential):
def forward(self, inputs : Tuple[Tensor, Tensor]):
def forward(self, inputs : Tuple[Tensor, Tuple[int, int]]):
for module in self._modules.values():
#if type(inputs) == tuple:
inputs = module(*inputs)

Loading…
Cancel
Save