Project Setup for Machine Learning with PyTorch and UV

While working on a recent project, I initially had some problems with setting up my python project with uv and adding pytorch as a dependency, so I thought I would quickly share my experience to save you from some headaches.

uv add torchvision
  × No solution found when resolving dependencies:
  ╰─▶ Because only the following versions of torchvision are available:
          torchvision==0.1.6
...
      and torchvision==0.1.6 was yanked (reason: 0.1.6 is past it's support date and confuses users on unsupported platforms), we can conclude that torchvision<0.1.7 cannot be used.
      And because torchvision>=0.1.7,<=0.2.2.post3 was yanked (reason: So that users won't accidentally install this when using python 3.11), we can conclude that torchvision<0.1.8 cannot be used.
      And because torchvision>=0.3.0,<=0.10.1 has no wheels with a matching Python version tag and torchvision==0.11.0 was yanked (reason: Dependency issue, depends on a version of torch that does not exist on pypi), we can conclude that torchvision<0.11.1 cannot be
      used.
      And because torchvision>=0.11.1,<=0.14.1 has no wheels with a matching Python version tag and torchvision==0.15.0 was yanked (reason: Contains an incorrect torch dependency), we can conclude that torchvision<0.15.1 cannot be used.
      And because torchvision>=0.15.1 has no wheels with a matching Python version tag and your project depends on torchvision, we can conclude that your project's requirements are unsatisfiable.
  help: If you want to add the package regardless of the failed resolution, provide the `--frozen` flag to skip locking and syncing.

To get started, let’s initialize a new project with uv. Make sure to specify the python version to be 3.12, as the latest python version 3.13 is not yet supported by pytorch at the time of writing this.

uv init my-awesome-ml-project --python 3.12

Now navigate into the project directory, and create a python virtual environment with the following command.

uv venv

PyTorch offers different installation variants, depending on whether you want to run your code on the GPU with Cuda or on your CPU. Depending on which variant you want, we need to select the appropriate source url and specify it in the pyproject.toml file. For example for cuda 11.8, you can add the following section to the project configuration file.

[[tool.uv.index]]
name = "pytorch-cu118"
url = "https://download.pytorch.org/whl/cu118"
explicit = true

Now we can install torch with uv. I also added numpy, for me it didn’t install it automatically, and I got a warning when importing torch without it.

uv add torch numpy

Activate the project virtual environment and run a quick test to see if you can import torch and check if cuda is installed.

python
>>> import torch
>>> torch.cuda.is_available()
True

Conclusion

In this short guide, I showed you how to setup your project using uv with pytorch. For more details, check out the guide on the UV website.

Happy coding!