Hi community, I have the following issue :
Following this memory leak issue Torchvision decode_jpeg memory leak · Issue #4378 · pytorch/vision · GitHub that seems to be fixed by upgrading to cuda 11.8, I’m trying to bump my application to this version and run into the following issue:
I’m using Pex to build and ship the app, and it struggles to create the pex file because the torch2.0.1+cu118 is larger than 2GB: torch-2.0.1+cu118-cp39-cp39-linux_x86_64.whl is too big (2265684151 bytes) to fit in memory (I think the message is misleading here, as I have plenty available memory when building)
When using torch2.0.1+cu117, it works because it is not shipped as a big wheel, but as a torch wheel + nvidia wheels (nvidia_cuda_runtime_xxxx, nvidia_cublas_xxxxx, etc) and as such all individual wheels are much smaller so Pex has no issue to build the app.
Is there an alternative repository or a way to install torch2.0.1+cu18 with separate wheels for nvidia libraries ?