Custom built Pytorch in docker

[archie1983]

Hi,
I’ve got a Jetson XAVIER with JetPack 5.1.3 / L4T 35.5.0. Within my docker container I need a Pytorch version that has been built with the USE_DISTRIBUTED=1 key. Pytorch 2.1 or 2.2 would do with an appropriate torchvision. If I build it on my laptop and then install the wheel in the docker container and then run my code, then I get the following:

File “/usr/local/lib/python3.8/dist-packages/torch/_ops.py”, line 692, in call
return self._op(*args, **kwargs or {})
NotImplementedError: Could not run ‘torchvision::nms’ with arguments from the ‘CUDA’ backend. This could be because the operator doesn’t exist for this backend, or was omitted during the selective/custom build process (if using custom build).

Next step was to rebuild your container (preferebaly l4t-pytorch, but pytorch would be a good start). So I checked out jetson-containers and modified

packages/l4t/l4t-pytorch/l4t-pytorch.json

to include “build_args”: {“USE_DISTRIBUTED” : “1”}

Then I start the whole process with: PYTORCH_VERSION=2.1 jetson-containers build l4t-pytorch

but unfortunately it doesn’t even try to build pytorch. It downloads a pre-built wheel and installs that.

Is there a way to force the build with the build arg I need?

Thanks

[AastaLLL]

Hi,

You can modify the Dockerfile below to make it build from the source:

dusty-nv/jetson-containers/blob/master/packages/pytorch/Dockerfile#L38

Thanks