Compiling pytorch with CUDA11.1+CUDNN 8.0.5, NO AVX with Ampere architecture

Hello everyone,

I am trying to compile pytorch on an Ubuntu 20 machine that does not have AVX support. The GPU on the server is based on the Ampere architecture with SM_86. I know that Cuda 11.1 and CUDNN 8.0.5 have SM_86 support and are not dependent on AVX (I successfully compiled Tensorflow, for example, and it works with the hardware).

I am looking for a combination of pytorch version and its dependencies that I compile with no AVX+the aforementioned CUDA/CUDNN versions. I cloned the github (recursively) and tried compiling different releases with the following options:


I’ve tried both clang and gcc, but the compilation stops at different points with different release versions of pytorch. I get errors at pytorch-cpu or ATEN and sometimes one of the dependencies seems to assume that a later version of CUDA is present (judging by the header constants it is unable to find).

Does anyone have any suggestions?

Many thanks!