[bug] Can't import smdistributed from the pytorch-training:2.4.0-gpu-py311 container

[gui-miotto]

Summary

Can't import smdistributed in the new pytorch image

To reproduce

Running

import smdistributed.dataparallel.torch.torch_smddp

will result in an error:

ModuleNotFoundError: No module named 'smdistributed'

This happens when using the pytorch 2.4 image, more precisely, the 763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.4.0-gpu-py311-cu124-ubuntu22.04-sagemaker

smdistributed could be imported in the older version (pytorch 2.3 and 2.2)

[gui-miotto]

Related to this