AttributeError: module 'transformer_engine' has no attribute 'pytorch'

[carrot0117]

I encountered the error AttributeError: module 'transformer_engine' has no attribute 'pytorch' when installing TE using pip install transformer_engine. Does anyone have any good solutions?

[carrot0117]

transformer_engine==1.13.0
pytorch==2.3.0
cuda==12.1

[ptrendx]

Hi @carrot0117 to install from the pip wheel you need to specify which framework you want to use TE with. E.g. to install with pyTorch support you need to install

pip install transformer_engine[pytorch]

For more information see our Installation guide.