How to specify a version of pytorch for l4t-pytorch?

dusty-nv / jetson-containers

Machine Learning Containers for NVIDIA Jetson and JetPack-L4T

MIT License

1.93k stars 422 forks source link

How to specify a version of pytorch for l4t-pytorch? #413

Open h-fernand opened 4 months ago

h-fernand commented 4 months ago

I would like to either run or build the l4t-pytorch container using the 2.0-distributed version of pytorch. I am trying to run MMDetection on my Jetson which requires the distributed packages to be available even when running in a non-distributed environment (for reasons that completely escape me). Would I accomplish this by doing something like ./build.sh l4t-pytorch pytorch:2.0-distributed or is the process more involved?

dusty-nv commented 4 months ago

Hi @h-fernand, the ones for JP6 are already enabled with distributed by default, however if you're on JP5, then just build pytorch:distributed . Then base your container off that, or change the depends in other jetson-containers to point to pytorch-distributed instead

TrackAnything container does this, and installs mmcv: https://github.com/dusty-nv/jetson-containers/blob/ac51db233cf5cbc9345c451267dfe29b06bff614/packages/vit/tam/Dockerfile