back to cu121

Dockerfile

@@ -5,7 +5,7 @@
 # for development purposes.
 # ----------------------------------------------------------------------------#
 # ~ CUDA + cuDNN on Ubuntu ~ #
-FROM nvcr.io/nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
+FROM nvcr.io/nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04
 # Prevents Python from creating __pycache__/ and .pyc/ folders in the project
 # folder
 ENV PYTHONPYCACHEPREFIX=/.cache/python/

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.0+cu124", # Tensor manipulation
+    "torch==2.5.0+cu121", # 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.5.0+cu124", # Tensor manipulation on audio data
-    "torchvision==0.20.0+cu124", # Tensor manipulation on vision data
+    "torchaudio==2.5.0+cu121", # Tensor manipulation on audio data
+    "torchvision==0.20.0+cu121", # Tensor manipulation on vision data
     "transformers==4.45.2", # Pre-trained models published on Hugging Face
     "diffusers==0.30.3", # 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/cu124"
+url = "https://download.pytorch.org/whl/cu121"
 explicit = true
 
 [tool.black]