diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix index 4217b87e81e7..95d2486bfbb2 100644 --- a/pkgs/development/python-modules/torch/default.nix +++ b/pkgs/development/python-modules/torch/default.nix @@ -16,9 +16,10 @@ filelock, jinja2, networkx, - openai-triton, sympy, numpy, pyyaml, cffi, click, typing-extensions, + # ROCm build and `torch.compile` requires `openai-triton` + tritonSupport ? (!stdenv.isDarwin), openai-triton, # Unit tests hypothesis, psutil, @@ -377,12 +378,10 @@ in buildPythonPackage rec { # the following are required for tensorboard support pillow six future tensorboard protobuf - # ROCm build and `torch.compile` requires openai-triton - openai-triton - # torch/csrc requires `pybind11` at runtime pybind11 ] + ++ lib.optionals tritonSupport [ openai-triton ] ++ lib.optionals MPISupport [ mpi ] ++ lib.optionals rocmSupport [ rocmtoolkit_joined ];