From 54ca7ceeda1ebbb6b6ec2358984e53f28d711825 Mon Sep 17 00:00:00 2001
From: rainx0r <angelos.h@hotmail.com>
Date: Mon, 19 Aug 2024 19:58:28 +0100
Subject: [PATCH] Update container

---
 metaworld-jax/requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/metaworld-jax/requirements.txt b/metaworld-jax/requirements.txt
index d802543..d6f8cd6 100644
--- a/metaworld-jax/requirements.txt
+++ b/metaworld-jax/requirements.txt
@@ -1,5 +1,5 @@
 # Jax
-jax[cuda12_pip]==0.4.31
+jax[cuda12]==0.4.31
 flax==0.8.5
 distrax==0.1.5