Exporting UnetPlusPlus as TorchScript gives error

[atncai]

Hi,

I try to export my models to TorchScript. Unet and DeepLaapV3Plus work fine, but when trying it on UnetPlusPlus, I get a runtime error:

model = smp.UnetPlusPlus(encoder_name="efficientnet-b5", encoder_weights="imagenet", in_channels=3, classes=5)

scripted_model = torch.jit.script(model)

RuntimeError: Unable to extract string literal index. ModuleDict indexing is only supported with string literals. For example, 'i = "a"; self.layers[i](x)' will fail because i is not a literal. Enumeration of ModuleDict is supported, e.g. 'for k, v in self.items(): out = v(inp)': File "/home/christian/dev/HSP/libraries/atnc_mcv/venv/lib/python3.12/site-packages/segmentation_models_pytorch/decoders/unetplusplus/decoder.py", line 153 for depth_idx in range(self.depth - layer_idx): if layer_idx == 0: output = self.blocks[f"x_{depth_idx}_{depth_idx}"]( ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE features[depth_idx], features[depth_idx + 1] )

Thanks, Christian