Fix: Use CUDA 12.4 on Linux (#501)

Author: m-naumann

Dockerfile

@@ -1,5 +1,5 @@
 # Base image with CUDA and cuDNN support                                                                                      
-FROM nvidia/cuda:12.2.2-cudnn8-devel-ubuntu22.04
+FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
 
 ARG DEBIAN_FRONTEND=noninteractive
 

pyproject.toml

@@ -62,6 +62,23 @@ vbd = [
 default-groups = "all"
 cache-dir = "./.uv_cache"
 
+# Use the CUDA index for torch on Linux
+[[tool.uv.index]]
+name = "pytorch-cuda"
+url = "https://download.pytorch.org/whl/cu124"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[tool.uv.sources]
+torch = [
+    { index = "pytorch-cuda", marker = "platform_system == 'Linux'" },
+    { index = "pytorch-cpu", marker = "platform_system != 'Linux'" }
+]
+
 [tool.madrona.packages.madrona_gpudrive]
 ext-only = true
 ext-out-dir = "build"