fix(hpc): modules loading

env/hpc/modules.txt

@@ -1,8 +1,6 @@
-PyTorch/2.1.2-foss-2023a-CUDA-12.1.1
-jax/0.4.25-gfbf-2023a-CUDA-12.1.1
-Flax/0.8.4-gfbf-2023a-CUDA-12.1.1
-Optax/0.2.2-gfbf-2023a-CUDA-12.1.1
-wandb/0.16.1-GCC-12.3.0
-matplotlib/3.7.2-gfbf-2023a
-PyYAML/6.0-GCCcore-12.3.0
-FFmpeg/6.0-GCCcore-12.3.0
+gfbf/2024a
+GCCcore/13.3.0
+Python/3.12.3-GCCcore-13.3.0
+PyTorch/2.6.0-foss-2024a-CUDA-12.6.0
+FFmpeg/7.0.2-GCCcore-13.3.0
+PyYAML/6.0.2-GCCcore-13.3.0

env/hpc/requirements.txt

@@ -1,9 +1,14 @@
 biorobot==0.4.2
 cleanrl>=0.4.8
 evosax==0.2.0
+flax>=0.12.2
 gymnasium>=1.2.3
 ipykernel==7.2.0
+jax==0.9.0.1
+matplotlib==3.10.8
 mediapy==1.2.6
+optax>=0.2.6
 pyopengl>=3.1.10
 pyopengl-accelerate>=3.1.10
 tyro>=1.0.10
+wandb==0.24.2