Converting to TorchScript using scripting. #669
Description
Currently the models from this repository cannot be converted to TorchScript using scripting. Ie. the following fails:
import torch.jit
import segmentation_models_pytorch as smp
model = smp.Unet()
script = torch.jit.script(model)With this error:
RuntimeError:
tuple slice indices must be integer constants:
File "/tmp/2022-10-10-13-30-03/venv/lib/python3.10/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
Converting models using scripting has a huge benefit of preserving dynamic logic. I can look into making this conversion possible, but I thought I'd discuss this change first. Perhaps this has already been attempted in the past?
ps. if you opt-in for Hacktoberfest it also helps my progress there ;p