diff --git a/Dockerfile b/Dockerfile index 72bb2d30..d8886ec3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -5,7 +5,7 @@ # for development purposes. # ----------------------------------------------------------------------------# # ~ CUDA + cuDNN on Ubuntu ~ # -FROM nvcr.io/nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 +FROM nvcr.io/nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 # Prevents Python from creating __pycache__/ and .pyc/ folders in the project # folder ENV PYTHONPYCACHEPREFIX=/.cache/python/ diff --git a/pyproject.toml b/pyproject.toml index 40dae190..cdfa57ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -10,7 +10,7 @@ version = "0.0.1" # The following assumes that we build on top of the NGC PyTorch Docker image. dependencies = [ # MANDATORY for cneuromax/ - "torch==2.5.1", # Tensor manipulation + "torch==2.5.1+cu124", # Tensor manipulation "hydra-core==1.3.2", # Configuration management "hydra-zen==0.13.0", # Cleaner Hydra configs "beartype==0.19.0", # Dynamic type checking @@ -49,8 +49,8 @@ dependencies = [ # MANDATORY for cneuromax/fitting/deeplearning/ "lightning==2.4.0", # PyTorch wrapper to decrease boilerplate # OPTIONAL for cneuromax/fitting/deeplearning/ - "torchaudio==2.4.1+cu121", # Tensor manipulation on audio data - "torchvision==0.19.1+cu121", # Tensor manipulation on vision data + "torchaudio==2.5.1+cu124", # Tensor manipulation on audio data + "torchvision==0.20.1+cu124", # Tensor manipulation on vision data "transformers==4.46.0", # Pre-trained models published on Hugging Face "diffusers==0.31.0", # Diffusion models published on Hugging Face "timm==1.0.11", # Image models @@ -95,7 +95,7 @@ torchvision = { index = "pytorch" } [[tool.uv.index]] name = "pytorch" -url = "https://download.pytorch.org/whl/cu121" +url = "https://download.pytorch.org/whl/cu124" explicit = true [tool.black]