How to specify a nightly Pytorch version in environment.yml

Consider the following environment.yml file:

channels:
  - nvidia
  - defaults
  - conda-forge
dependencies:
  - bottleneck>=1.3.6
  - catboost>=0.24.4
  ...
  - pip:
    - xgboost==2.1.4
    ...

How do I add the following pip command to the yml file without disrupting the current behavior of the other pip install commands above?

pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

Depending on what exactly you mean by without disrupting the current behavior of the other pip install commands, the following may or may not be a viable solution for the - pip section of your environment.yml file:

channels:
  - nvidia
  - defaults
  - conda-forge
dependencies:
  - pip:
    - --pre
    - --extra-index-url https://download.pytorch.org/whl/nightly/cu128
    - xgboost==2.1.4
    - torch
    - torchvision
    - torchaudio

Why I think this should work in your case

• The --pre flag is a global option for pip requirements, enabling the installation of pre-release versions. However, since you explicitly specified the exact version of all other shown packages (xgboost==2.1.4, in this case), it will only apply to those without an exact version specified (torch, torchvision, torchaudio, in this case).
• Using --extra-index-url instead of --index-url (which are also both global options) will ensure that (a) packages not present in the specified URL will be installed from PyPI (i.e. from the default index), (b) packages present in the specified URL will be potentially installed from there, including pre-release versions since we specified the --pre flag.

Using the corresponding YAML file and installing it with conda leads to the following packages being installed via pip in my case:

filelock                  3.18.0                   pypi_0    pypi
fsspec                    2025.3.0                 pypi_0    pypi
jinja2                    3.1.6                    pypi_0    pypi
markupsafe                3.0.2                    pypi_0    pypi
mpmath                    1.3.0                    pypi_0    pypi
networkx                  3.4.2                    pypi_0    pypi
numpy                     2.2.3                    pypi_0    pypi
nvidia-cublas-cu12        12.8.3.14                pypi_0    pypi
nvidia-cuda-cupti-cu12    12.8.57                  pypi_0    pypi
nvidia-cuda-nvrtc-cu12    12.8.61                  pypi_0    pypi
nvidia-cuda-runtime-cu12  12.8.57                  pypi_0    pypi
nvidia-cudnn-cu12         9.7.1.26                 pypi_0    pypi
nvidia-cufft-cu12         11.3.3.41                pypi_0    pypi
nvidia-cufile-cu12        1.13.0.11                pypi_0    pypi
nvidia-curand-cu12        10.3.9.55                pypi_0    pypi
nvidia-cusolver-cu12      11.7.2.55                pypi_0    pypi
nvidia-cusparse-cu12      12.5.7.53                pypi_0    pypi
nvidia-cusparselt-cu12    0.6.3                    pypi_0    pypi
nvidia-nccl-cu12          2.25.1                   pypi_0    pypi
nvidia-nvjitlink-cu12     12.8.61                  pypi_0    pypi
nvidia-nvtx-cu12          12.8.55                  pypi_0    pypi
pillow                    11.1.0                   pypi_0    pypi
pytorch-triton            3.3.0+git96316ce5          pypi_0    pypi
scipy                     1.15.2                   pypi_0    pypi
sympy                     1.13.3                   pypi_0    pypi
torch                     2.8.0.dev20250313+cu128          pypi_0    pypi
torchaudio                2.6.0.dev20250313+cu128          pypi_0    pypi
torchvision               0.22.0.dev20250313+cu128          pypi_0    pypi
typing-extensions         4.12.2                   pypi_0    pypi
xgboost                   2.1.4                    pypi_0    pypi

So, the only pre-release versions are indeed the ones for the PyTorch packages and their dependencies.

Why this solution might still be problematic

As noted above, the --pre flag is a global option. So if by without disrupting the current behavior of the other pip install commands, you mean that the --pre flag should only be applied to the PyTorch packages and their potential dependencies, this won't work unless all other packages and their dependencies are specified explicitly in a way that does not allow for pre-release versions.

Alternatively you could, of course, specify the exact pre-release version of the PyTorch packages and dump the --pre flag, but I guess that would defy the purpose of what you currently achieve with your pip install call (namely, installing the latest nightlies without exactly specifying them).

What did not work

I tried some other approaches that turned out to fail:

Leaving out the --pre flag entirely

You might have noticed that the call …

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu128

… will still install the pre-release versions of torch, torchvision, and torchaudio, even though the --pre flag is omitted. This may appear counter-intuitive at first, since the pip docs explicitly state:

Starting with v1.4, pip will only install stable versions as specified by pre-releases by default.

However if we follow the link from the citation, we will find there:

Pre-releases of any kind … are implicitly excluded from all version specifiers, unless … the only available version that satisfies the version specifier is a pre-release.

This means:

• If we only specify --index-url https://download.pytorch.org/whl/nightly/cu128, then the pre-release versions of the PyTorch packages will be installed even without the --pre flag, because under the given URL, only pre-releases can be found, so each requirement can only be satisfied by installing a pre-release.

  The problem here: Packages that cannot be found under this URL (xgboost in the given case) cannot be installed at all.

• If we specify --extra-index-url https://download.pytorch.org/whl/nightly/cu128 instead, and don't provide the --pre flag, then xgboost can be installed.

  The problem here: Now release versions of the PyTorch packages will be installed, since they can be found on PyPI and thus the pre-release versions are no longer the only option.

Providing two separate - pip sections in the environment.yml file

What conda does when installing the pip dependencies, is essentially (1) creating a temporary requirements.txt file from the - pip section of the YAML file, then (2) installing this requirements.txt file via a pip install -r … call. What I thought could work would be creating two separate - pip sections – one for your "regular" pip dependencies and one for the PyTorch packages. I thought that then perhaps pip install -r … would be called separately for each section.

However, all but the last such - pip section in the YAML file seem to be ignored.

Referencing an external requirements.txt file

I also tried to specify the following torch_requirements.txt file:

--index-url https://download.pytorch.org/whl/nightly/cu128
torch
torchvision
torchaudio

This should work on its own (in particular, it should install the pre-release versions even without specifying the --pre flag; see above). I then tried to include this in the YAML file as:

...
dependencies:
  - pip:
    - xgboost==2.1.4
    - -r torch_requirements.txt

However, while working in principle (i.e. incorporating the contents of the referenced file), this again made the --index-url apply globally, so prevented to install packages like xgboost from PyPI, as described above.

Bottom line

With the way that conda currently handles pip packages, I don't think there is a universally applicable solution. The particular problem, in my opinion, is that there is no way (at least I did not find any) to trigger more than one separate pip install … call. The latter, however, would be necessary to apply different global pip requirement options (such as --pre or --index-url) for different sets of pip packages.