diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index 7a399f54512a..d889dfc60201 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -1,6 +1,6 @@ { stdenv, lib, fetchFromGitHub, fetchpatch, buildPythonPackage, python, cudaSupport ? false, cudaPackages, magma, - mklDnnSupport ? true, useSystemNccl ? true, + useSystemNccl ? true, MPISupport ? false, mpi, buildDocs ? false, @@ -17,6 +17,10 @@ # Unit tests hypothesis, psutil, + # Disable MKLDNN on aarch64-darwin, it negatively impacts performance, + # this is also what official pytorch build does + mklDnnSupport ? !(stdenv.isDarwin && stdenv.isAarch64), + # virtual pkg that consistently instantiates blas across nixpkgs # See https://github.com/NixOS/nixpkgs/pull/83888 blas,