I've trained a segmentation model using Python 3.8 environment and segmentation_models_pytorch aka smp. When I saved it and load in my prediction environment (Python 3.6 with smp) it worked with just

import torch

model = torch.load(path.join('models', model_name))

However it conflicts with onnx package (onnx requires newer Python). I've created new conda environment with Python 3.10 (another with Python 3.11). Now torch refuses to load the model with an error message ModuleNotFoundError: No module named 'segmentation_models_pytorch.unet'.

What is the right way to

  1. install torch that supports torch.load

  2. torch that knows about unet architecture

  3. on the environment that supports onnx model conversion

  4. doesnot require some old and very special version of Python (which causes conflicts with other packages)

    from os import path, environ
    import torch
    torch_model = torch.load(path.join('models', 'my_model.pth'))
    Traceback (most recent call last):
    File "D:\workspace\acne_prod\pytorch2onnx.py", line 3, in <module>
    torch_model = torch_load(path.join('models', 'test_model.pth'))
    File "C:\Users\sixty\anaconda3\envs\acne_prod_smp_onnx\Lib\site-packages\torch\serialization.py", line 809, in load
      return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
    File "C:\Users\sixty\anaconda3\envs\acne_prod_smp_onnx\Lib\site-packages\torch\serialization.py", line 1172, in _load
    result = unpickler.load()
    File "C:\Users\sixty\anaconda3\envs\acne_prod_smp_onnx\Lib\site-packages\torch\serialization.py", line 1165, in find_class
      return super().find_class(mod_name, name)
    ModuleNotFoundError: No module named 'segmentation_models_pytorch.unet'

2 Answers 2


Segmentation_models is a different python package from pytorch and requires separate installation. !pip install -U git+https://github.com/qubvel/segmentation_models.pytorch works for me on python 3.10.5 and later.

  • $\begingroup$ This is the first troubleshooting step. It didnot help in my case. See my solution below. $\endgroup$
    – sixtytrees
    Jun 1, 2023 at 19:11

Here is the solution that worked for me.

Downgrade your SMP to the version you were training at. Or train your model on a fresh SMP if resourses allow. This way you avoid legacy Python issues.

Hence. This is NOT a pytorch issue. This is due to the change of the SMP models that caused issues with dackward compatibility.


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.