From 3080e7c051d23df908bf17633db56b75f389b594 Mon Sep 17 00:00:00 2001
From: Samuel Ainsworth <skainsworth@gmail.com>
Date: Fri, 21 Jul 2023 13:31:30 -0700
Subject: [PATCH] python3Packages.torchvision-bin: fix aarch64-darwin build

---
 pkgs/development/python-modules/torchvision/bin.nix | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/pkgs/development/python-modules/torchvision/bin.nix b/pkgs/development/python-modules/torchvision/bin.nix
index bc33961c097..e1ec5367114 100644
--- a/pkgs/development/python-modules/torchvision/bin.nix
+++ b/pkgs/development/python-modules/torchvision/bin.nix
@@ -28,13 +28,14 @@ in buildPythonPackage {
 
   disabled = (pythonOlder "3.8") || (pythonAtLeast "3.12");
 
-  buildInputs = with cudaPackages; [
+  # Note that we don't rely on config.cudaSupport here, because the Linux wheels all come built with CUDA support.
+  buildInputs = with cudaPackages; lib.optionals stdenv.isLinux [
     # $out/${sitePackages}/torchvision/_C.so wants libcudart.so.11.0 but torchvision.libs only ships
     # libcudart.$hash.so.11.0
     cuda_cudart
   ];
 
-  nativeBuildInputs = [
+  nativeBuildInputs = lib.optionals stdenv.isLinux [
     autoPatchelfHook
     addOpenGLRunpath
   ];
@@ -49,7 +50,7 @@ in buildPythonPackage {
 
   pythonImportsCheck = [ "torchvision" ];
 
-  preInstall = ''
+  preInstall = lib.optionalString stdenv.isLinux ''
     addAutoPatchelfSearchPath "${torch-bin}/${python.sitePackages}/torch"
   '';
 
@@ -62,7 +63,7 @@ in buildPythonPackage {
     # https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
     license = licenses.bsd3;
     sourceProvenance = with sourceTypes; [ binaryNativeCode ];
-    platforms = [ "x86_64-linux" ];
+    platforms = [ "aarch64-darwin" "x86_64-darwin" "x86_64-linux" ];
     maintainers = with maintainers; [ junjihashimoto ];
   };
 }