python3Packages.pytorch3d: cleanup, fix build with cudaSupport

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
     ];
   };
-}
+})