diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index b00146b573fe..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, @@ -9,7 +9,7 @@ # Build inputs numactl, - CoreServices, libobjc, + Accelerate, CoreServices, libobjc, # Propagated build inputs numpy, pyyaml, cffi, click, typing-extensions, @@ -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, @@ -275,7 +279,7 @@ in buildPythonPackage rec { ++ lib.optionals rocmSupport [ openmp ] ++ lib.optionals (cudaSupport || rocmSupport) [ magma ] ++ lib.optionals stdenv.isLinux [ numactl ] - ++ lib.optionals stdenv.isDarwin [ CoreServices libobjc ]; + ++ lib.optionals stdenv.isDarwin [ Accelerate CoreServices libobjc ]; propagatedBuildInputs = [ cffi diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index c63ac709b295..e39f39cae322 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -11715,7 +11715,7 @@ self: super: with self; { torch = callPackage ../development/python-modules/torch { cudaSupport = pkgs.config.cudaSupport or false; - inherit (pkgs.darwin.apple_sdk.frameworks) CoreServices; + inherit (pkgs.darwin.apple_sdk.frameworks) Accelerate CoreServices; inherit (pkgs.darwin) libobjc; inherit (pkgs.llvmPackages_rocm) openmp; };