diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix index 80c3c6d7293..a746d4c7dfb 100644 --- a/pkgs/development/python-modules/pytorch/default.nix +++ b/pkgs/development/python-modules/pytorch/default.nix @@ -104,7 +104,7 @@ let "LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH "; in buildPythonPackage rec { - version = "1.5.0"; + version = "1.5.1"; pname = "pytorch"; disabled = !isPy3k; @@ -118,9 +118,24 @@ in buildPythonPackage rec { repo = "pytorch"; rev = "v${version}"; fetchSubmodules = true; - sha256 = "19qyrjd72mc0llcfn50av8ym05f2iwa38gv068wykji4ph7qjlv2"; + sha256 = "1xjbn4hi96m7xslv3p2jc6qcsng0fx3w1m6isqfah81piljf8wng"; }; + patches = lib.optionals stdenv.isAarch64 [ + # GNU aarch64 assembler does not support 4s on neon mov: + # https://github.com/pytorch/pytorch/issues/33124 + # + # Fix from: + # https://github.com/pytorch/pytorch/pull/40584 + # + # This patch can be removed with the next major version (1.6.0). + (fetchpatch { + name = "qnnpack-neon-fix.patch"; + url = "https://github.com/pytorch/pytorch/commit/7676682584d0caf9243bce74ea0a88711ec4a807.diff"; + sha256 = "13spncaqlpsp8qk2850yly7xqwmhhfwznhmzkk8jgpslkbx75vgq"; + }) + ]; + preConfigure = lib.optionalString cudaSupport '' export TORCH_CUDA_ARCH_LIST="${lib.strings.concatStringsSep ";" final_cudaArchList}" export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++