Skip to content

Commit c1d8056

Browse files
wylimilesial
authored andcommitted
ignore type hints in jit
Signed-off-by: Wenqi Li <wenqil@nvidia.com>
1 parent e2841b3 commit c1d8056

File tree

1 file changed

+3
-4
lines changed

1 file changed

+3
-4
lines changed

monai/networks/nets/dints.py

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import torch
1919
import torch.nn as nn
2020
import torch.nn.functional as F
21+
2122
from monai.networks.blocks.dints_block import (
2223
ActiConvNormBlock,
2324
FactorizedIncreaseBlock,
@@ -39,7 +40,7 @@
3940
class CellInterface(torch.nn.Module):
4041
"""interface for torchscriptable Cell"""
4142

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

4546

@@ -675,9 +676,7 @@ def forward(self, x: list[torch.Tensor]) -> list[torch.Tensor]:
675676
for res_idx, activation in enumerate(self.arch_code_a[blk_idx].data):
676677
if activation:
677678
mod: CellInterface = self.cell_tree[str((blk_idx, res_idx))]
678-
_out = mod.forward(
679-
x=inputs[self.arch_code2in[res_idx]], weight=None
680-
)
679+
_out = mod.forward(x=inputs[self.arch_code2in[res_idx]], weight=None)
681680
outputs[self.arch_code2out[res_idx]] = outputs[self.arch_code2out[res_idx]] + _out
682681
inputs = outputs
683682

0 commit comments

Comments
 (0)