Commit 8cee6dfd authored by Gaetan Lepage's avatar Gaetan Lepage
Browse files

python3Packages.pytorch3d: cleanup, fix build with cudaSupport

parent 9502c8fc

pkgs/development/python-modules/pytorch3d/default.nix

+27 −10
Original line number Diff line number Diff line
@@ -2,18 +2,23 @@ 
  lib,

  buildPythonPackage,

  fetchFromGitHub,

  cudaPackages,



  # build-system

  ninja,

  setuptools,

  wheel,

  torch,



  # dependencies

  iopath,

  cudaPackages,

  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

    torch

  ];



  dependencies = [

    torch

    iopath

    torch

  ];

  buildInputs = [

    (lib.getOutput "cxxdev" 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

    ];

  };

}
 
})