tuple slice indices must be integer constants error #755
Description
When trying to convert Unet segmentation model to torchscript model, get a run time error in _base.py that tuple slice indices must be integer constants.
config["model_config"] = {
"encoder_name" : "mobilenet_v2",
"encoder_weights" : "imagenet",
"encoder_depth" : 5,
"activation" : None,
"classes" : 3,
"decoder_channels" : [256, 128, 64, 32, 16]
}
model = smp.Unet(**config["model_config"])
On converting the model to torchscript using : traced_script_module = torch.jit.script(model), I am getting the following error.
---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
Cell In[15], line 6
3 sample = sample.cuda()
5 # traced_script_module = torch.jit.trace(model, sample)
----> 6 traced_script_module = torch.jit.script(model)
8 traced_script_module.save(os.path.join(model_path, "fork-detection.v2.pt"))
9 traced_script_module.save(os.path.join(model_path, "fork-detection.v2.ts"))
File ~/miniconda3/envs/fork/lib/python3.9/site-packages/torch/jit/_script.py:1286, in script(obj, optimize, _frames_up, _rcb, example_inputs)
1284 if isinstance(obj, torch.nn.Module):
1285 obj = call_prepare_scriptable_func(obj)
-> 1286 return torch.jit._recursive.create_script_module(
1287 obj, torch.jit._recursive.infer_methods_to_compile
1288 )
1290 if isinstance(obj, dict):
1291 return create_script_dict(obj)
File ~/miniconda3/envs/fork/lib/python3.9/site-packages/torch/jit/_recursive.py:458, in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
456 if not is_tracing:
457 AttributeTypeIsSupportedChecker().check(nn_module)
--> 458 return create_script_module_impl(nn_module, concrete_type, stubs_fn)
File ~/miniconda3/envs/fork/lib/python3.9/site-packages/torch/jit/_recursive.py:520, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
517 script_module._concrete_type = concrete_type
519 # Actually create the ScriptModule, initializing it with the function we just defined
--> 520 script_module = torch.jit.RecursiveScriptModule._construct(cpp_module, init_fn)
522 # Compile methods if necessary
523 if concrete_type not in concrete_type_store.methods_compiled:
File ~/miniconda3/envs/fork/lib/python3.9/site-packages/torch/jit/_script.py:615, in RecursiveScriptModule._construct(cpp_module, init_fn)
602 """
603 Construct a RecursiveScriptModule that's ready for use. PyTorch
604 code should use this to construct a RecursiveScriptModule instead
(...)
612 init_fn: Lambda that initializes the RecursiveScriptModule passed to it.
613 """
614 script_module = RecursiveScriptModule(cpp_module)
--> 615 init_fn(script_module)
617 # Finalize the ScriptModule: replace the nn.Module state with our
618 # custom implementations and flip the _initializing bit.
619 RecursiveScriptModule._finalize_scriptmodule(script_module)
File ~/miniconda3/envs/fork/lib/python3.9/site-packages/torch/jit/_recursive.py:498, in create_script_module_impl.<locals>.init_fn(script_module)
495 scripted = orig_value
496 else:
497 # always reuse the provided stubs_fn to infer the methods to compile
--> 498 scripted = create_script_module_impl(orig_value, sub_concrete_type, stubs_fn)
500 cpp_module.setattr(name, scripted)
501 script_module._modules[name] = scripted
File ~/miniconda3/envs/fork/lib/python3.9/site-packages/torch/jit/_recursive.py:524, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
522 # Compile methods if necessary
523 if concrete_type not in concrete_type_store.methods_compiled:
--> 524 create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
525 # Create hooks after methods to ensure no name collisions between hooks and methods.
526 # If done before, hooks can overshadow methods that aren't exported.
527 create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs)
File ~/miniconda3/envs/fork/lib/python3.9/site-packages/torch/jit/_recursive.py:375, in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
372 property_defs = [p.def_ for p in property_stubs]
373 property_rcbs = [p.resolution_callback for p in property_stubs]
--> 375 concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)
RuntimeError:
tuple slice indices must be integer constants:
File "/home/chinmay/miniconda3/envs/fork/lib/python3.9/site-packages/segmentation_models_pytorch/encoders/_base.py", line 20
def out_channels(self):
"""Return channels dimensions for each tensor of forward output of encoder"""
return self._out_channels[: self._depth + 1]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE