Unverified Commit 585ce2c4 authored by F-G Fernandez's avatar F-G Fernandez Committed by GitHub
Browse files

Fixed typing exception throwing issues with JIT (#3029)

* refactor: Fixed typing exception throwing issues with JIT

* style: Added back mypy typing to forward

* chore: Added back densenet module as mypy exception
parent 9c799349
...@@ -12,6 +12,10 @@ ignore_errors = True ...@@ -12,6 +12,10 @@ ignore_errors = True
ignore_errors = True ignore_errors = True
[mypy-torchvision.models.densenet.*]
ignore_errors=True
[mypy-torchvision.models.detection.*] [mypy-torchvision.models.detection.*]
ignore_errors = True ignore_errors = True
......
...@@ -71,13 +71,13 @@ class _DenseLayer(nn.Module): ...@@ -71,13 +71,13 @@ class _DenseLayer(nn.Module):
def forward(self, input: List[Tensor]) -> Tensor: def forward(self, input: List[Tensor]) -> Tensor:
pass pass
@torch.jit._overload_method # type: ignore[no-redef] # noqa: F811 @torch.jit._overload_method # noqa: F811
def forward(self, input: Tensor) -> Tensor: def forward(self, input: Tensor) -> Tensor:
pass pass
# torchscript does not yet support *args, so we overload method # torchscript does not yet support *args, so we overload method
# allowing it to take either a List[Tensor] or single Tensor # allowing it to take either a List[Tensor] or single Tensor
def forward(self, input: Tensor) -> Tensor: # type: ignore[no-redef] # noqa: F811 def forward(self, input: Tensor) -> Tensor: # noqa: F811
if isinstance(input, Tensor): if isinstance(input, Tensor):
prev_features = [input] prev_features = [input]
else: else:
...@@ -121,7 +121,7 @@ class _DenseBlock(nn.ModuleDict): ...@@ -121,7 +121,7 @@ class _DenseBlock(nn.ModuleDict):
) )
self.add_module('denselayer%d' % (i + 1), layer) self.add_module('denselayer%d' % (i + 1), layer)
def forward(self, init_features: Tensor) -> Tensor: # type: ignore[override] def forward(self, init_features: Tensor) -> Tensor:
features = [init_features] features = [init_features]
for name, layer in self.items(): for name, layer in self.items():
new_features = layer(features) new_features = layer(features)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment