@Morganh
Here’s the content of the dockerfile I used for the environment:
FROM nvcr.io/nvidia/l4t-tensorrt:r8.5.2.2-devel
RUN apt update
RUN apt --fix-broken install
RUN apt install -y mpich
RUN pip install mpi4py
RUN pip install torchinfo
RUN pip install clearml
RUN pip install segmentation-models-pytorch
RUN pip install transformers
RUN pip install nvidia-tao-deploy