Update davit.py

pull/1630/head
Fredo Guan 3 years ago
parent 5c73dd66cb
commit 6f1cfd1aa2

@ -45,10 +45,10 @@ class SequentialWithSize(nn.Sequential):
class SequentialWithSize(nn.Sequential): class SequentialWithSize(nn.Sequential):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
super(SequentialWithSize, nn.Sequential, self).__init__(*args, **kwargs) super().__init__(*args, **kwargs)
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 _modules.values():
x, size = module(x, size) x, size = module(x, size)
return x, size return x, size

Loading…
Cancel
Save