diff --git a/pkgs/development/python-modules/pytorch3d/default.nix b/pkgs/development/python-modules/pytorch3d/default.nix new file mode 100644 index 000000000000..1a9151591863 --- /dev/null +++ b/pkgs/development/python-modules/pytorch3d/default.nix @@ -0,0 +1,58 @@ +{ + lib, + buildPythonPackage, + fetchFromGitHub, + setuptools, + wheel, + torch, + iopath, + cudaPackages, + config, + cudaSupport ? config.cudaSupport, +}: + +assert cudaSupport -> torch.cudaSupport; + +buildPythonPackage rec { + pname = "pytorch3d"; + version = "0.7.8"; + pyproject = true; + + src = fetchFromGitHub { + owner = "facebookresearch"; + repo = "pytorch3d"; + rev = "V${version}"; + hash = "sha256-DEEWWfjwjuXGc0WQInDTmtnWSIDUifyByxdg7hpdHlo="; + }; + + nativeBuildInputs = lib.optionals cudaSupport [ cudaPackages.cuda_nvcc ]; + build-system = [ + setuptools + wheel + ]; + dependencies = [ + torch + iopath + ]; + buildInputs = [ (lib.getOutput "cxxdev" torch) ]; + + env = + { + FORCE_CUDA = cudaSupport; + } + // lib.optionalAttrs cudaSupport { + TORCH_CUDA_ARCH_LIST = "${lib.concatStringsSep ";" torch.cudaCapabilities}"; + }; + + pythonImportsCheck = [ "pytorch3d" ]; + + meta = { + description = "FAIR's library of reusable components for deep learning with 3D data"; + homepage = "https://github.com/facebookresearch/pytorch3d"; + license = lib.licenses.bsd3; + maintainers = with lib.maintainers; [ + pbsds + SomeoneSerge + ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index c09180a5bd4b..8d5c14269ca5 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -10467,6 +10467,8 @@ self: super: with self; { pytomorrowio = callPackage ../development/python-modules/pytomorrowio { }; + pytorch3d = callPackage ../development/python-modules/pytorch3d { }; + pytouchlinesl = callPackage ../development/python-modules/pytouchlinesl { }; pyuca = callPackage ../development/python-modules/pyuca { };