tuple slice indices must be integer constants error

[nefario7]

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

[UCPRER]

Hello, I have encountered the same problem, have you resolved it ?

[eikaramba]

i have the same problem. would be great to have more infos or a solution here

[eikaramba]

ok got it. you can trace instead of script.

x = torch.randn(1, 3, 256, 256)
#scripted_module = torch.jit.script(model)
scripted_module = torch.jit.trace(model,x)

[diya97va]

I was trying to convert Unet segmentation model to torchscript model and ran into same error. Please let me know whether script worked for you rather than trace @nefario7