Skip to content

Commit

Permalink
ignore type hints in jit
Browse files Browse the repository at this point in the history
Signed-off-by: Wenqi Li <wenqil@nvidia.com>
  • Loading branch information
wyli authored and milesial committed Feb 23, 2023
1 parent e2841b3 commit c1d8056
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions monai/networks/nets/dints.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
import torch
import torch.nn as nn
import torch.nn.functional as F

from monai.networks.blocks.dints_block import (
ActiConvNormBlock,
FactorizedIncreaseBlock,
Expand All @@ -39,7 +40,7 @@
class CellInterface(torch.nn.Module):
"""interface for torchscriptable Cell"""

def forward(self, x: torch.Tensor, weight: torch.Tensor | None) -> torch.Tensor: # type: ignore
def forward(self, x: torch.Tensor, weight) -> torch.Tensor: # type: ignore
pass


Expand Down Expand Up @@ -675,9 +676,7 @@ def forward(self, x: list[torch.Tensor]) -> list[torch.Tensor]:
for res_idx, activation in enumerate(self.arch_code_a[blk_idx].data):
if activation:
mod: CellInterface = self.cell_tree[str((blk_idx, res_idx))]
_out = mod.forward(
x=inputs[self.arch_code2in[res_idx]], weight=None
)
_out = mod.forward(x=inputs[self.arch_code2in[res_idx]], weight=None)
outputs[self.arch_code2out[res_idx]] = outputs[self.arch_code2out[res_idx]] + _out
inputs = outputs

Expand Down

0 comments on commit c1d8056

Please sign in to comment.