DiceLoss can't support TorchScript #1592
Description
The DiceLoss reports error when torch.jit.script():
Traceback (most recent call last):
File "/home/hroth/Code2/pytorch-pipeline_covid-fl/baseline_supernet/supernet_fitter.py", line 50, in __init__
self.criterion = torch.jit.script(criterion)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_script.py", line 875, in script
obj, torch.jit._recursive.infer_methods_to_compile
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 317, in create_script_module
return create_script_module_impl(nn_module, concrete_type, stubs_fn)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 372, in create_script_module_impl
script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_script.py", line 365, in _construct
init_fn(script_module)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 352, in init_fn
scripted = create_script_module_impl(orig_value, sub_concrete_type, infer_methods_to_compile)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 376, in create_script_module_impl
create_methods_from_stubs(concrete_type, stubs)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 292, in create_methods_from_stubs
concrete_type._create_methods(defs, rcbs, defaults)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 596, in try_compile_fn
return torch.jit.script(fn, _rcb=rcb)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_script.py", line 917, in script
qualified_name, ast, _rcb, get_default_args(obj)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_recursive.py", line 596, in try_compile_fn
return torch.jit.script(fn, _rcb=rcb)
File "/opt/conda/lib/python3.6/site-packages/torch/jit/_script.py", line 917, in script
qualified_name, ast, _rcb, get_default_args(obj)
RuntimeError:
Expression of type dots cannot be used in a type expression:
File "/opt/monai/monai/utils/misc.py", line 70
def ensure_tuple_size(tup: Any, dim: int, pad_val: Any = 0) -> Tuple[Any, ...]:
~~~ <--- HERE
"""
Returns a copy of `tup` with `dim` values by either shortened or padded with `pad_val` as necessary.
'ensure_tuple_size' is being compiled since it was called from 'one_hot'
File "/opt/monai/monai/networks/utils.py", line 49
# if `dim` is bigger, add singleton dim at the end
if labels.ndim < dim + 1:
shape = ensure_tuple_size(labels.shape, dim + 1, 1)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
labels = labels.reshape(*shape)
'one_hot' is being compiled since it was called from 'DiceLoss.forward'
File "/opt/monai/monai/losses/dice.py", line 127
warnings.warn("single channel prediction, `to_onehot_y=True` ignored.")
else:
target = one_hot(target, num_classes=n_pred_ch)
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
if not self.include_background: