From 8cee6dfd2254761ca002eb4a4781ef66ce488b01 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 15 Apr 2026 13:29:36 +0000
Subject: [PATCH] python3Packages.pytorch3d: cleanup, fix build with
 cudaSupport

---
 .../python-modules/pytorch3d/default.nix      | 43 +++++++++++++------
 1 file changed, 30 insertions(+), 13 deletions(-)

diff --git a/pkgs/development/python-modules/pytorch3d/default.nix b/pkgs/development/python-modules/pytorch3d/default.nix
index 98896630ae81..40dbfd754a58 100644
--- a/pkgs/development/python-modules/pytorch3d/default.nix
+++ b/pkgs/development/python-modules/pytorch3d/default.nix
@@ -2,18 +2,23 @@
   lib,
   buildPythonPackage,
   fetchFromGitHub,
-  setuptools,
-  wheel,
-  torch,
-  iopath,
   cudaPackages,
+
+  # build-system
+  ninja,
+  setuptools,
+
+  # dependencies
+  iopath,
+  torch,
+
   config,
   cudaSupport ? config.cudaSupport,
 }:
 
 assert cudaSupport -> torch.cudaSupport;
 
-buildPythonPackage rec {
+buildPythonPackage.override { inherit (torch) stdenv; } (finalAttrs: {
   pname = "pytorch3d";
   version = "0.7.9";
   pyproject = true;
@@ -21,20 +26,27 @@ buildPythonPackage rec {
   src = fetchFromGitHub {
     owner = "facebookresearch";
     repo = "pytorch3d";
-    tag = "v${version}";
+    tag = "v${finalAttrs.version}";
     hash = "sha256-k+hqAy0T5ZReJ8jWuKJ/VTzjqd4qo8rrcFo8y4LJhhY=";
   };
 
-  nativeBuildInputs = lib.optionals cudaSupport [ cudaPackages.cuda_nvcc ];
+  nativeBuildInputs = lib.optionals cudaSupport [
+    cudaPackages.cuda_nvcc
+  ];
+
   build-system = [
+    ninja
     setuptools
-    wheel
-  ];
-  dependencies = [
     torch
-    iopath
   ];
-  buildInputs = [ (lib.getOutput "cxxdev" torch) ];
+
+  dependencies = [
+    iopath
+    torch
+  ];
+  buildInputs = [
+    (lib.getOutput "cxxdev" torch)
+  ];
 
   env = {
     FORCE_CUDA = cudaSupport;
@@ -43,6 +55,10 @@ buildPythonPackage rec {
     TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}";
   };
 
+  preBuild = ''
+    export MAX_JOBS=$NIX_BUILD_CORES
+  '';
+
   pythonImportsCheck = [ "pytorch3d" ];
 
   passthru.tests.rotations-cuda =
@@ -61,10 +77,11 @@ buildPythonPackage rec {
   meta = {
     description = "FAIR's library of reusable components for deep learning with 3D data";
     homepage = "https://github.com/facebookresearch/pytorch3d";
+    changelog = "https://github.com/facebookresearch/pytorch3d/releases/tag/${finalAttrs.src.tag}";
     license = lib.licenses.bsd3;
     maintainers = with lib.maintainers; [
       pbsds
       SomeoneSerge
     ];
   };
-}
+})