From 0f46e8359d087b03a9378614ff2df1704553ac40 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Fri, 27 Mar 2026 20:33:45 +0100 Subject: [PATCH] feat: switch managed venv to jax[cuda13] for GPU feature extraction RTX 6000 Pro (Blackwell SM 10.0) fully supports CUDA 13. Switch from jax[cpu]+jaxlib to jax[cuda13] which bundles jaxlib and uses pip-managed CUDA libraries. Delete _extract_env to force a rebuild. Co-Authored-By: Claude Sonnet 4.6 --- nodes/feature_extractor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nodes/feature_extractor.py b/nodes/feature_extractor.py index 16a4f9d..f15536c 100644 --- a/nodes/feature_extractor.py +++ b/nodes/feature_extractor.py @@ -17,7 +17,8 @@ _EXTRACT_PACKAGES = [ "torch", "torchaudio", "torchvision", # TF 2.15 only supports Python <=3.11; use >=2.16 for Python 3.12+ "tensorflow-cpu>=2.16.0", - "jax[cpu]", "jaxlib", + # jax[cuda13] includes jaxlib; pip-managed CUDA libs (no local toolkit needed) + "jax[cuda13]", "transformers", "decord", "einops", "numpy", "mediapy", "git+https://github.com/google-deepmind/videoprism.git", ]