add safetensors as dep for wrappers file

FLAIROx · Oct 14, 2024 · 1c81f41 · 1c81f41
1 parent 9c310f4
commit 1c81f41
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -34,6 +34,7 @@ dependencies = [
   "jax>=0.4.16.0,<=0.4.25",
   "jaxlib>=0.4.16.0,<=0.4.25",
   "flax",
+  "safetensors",
   "chex",
   "brax==0.10.3",
   "mujoco==3.1.3",
@@ -47,7 +48,6 @@ dependencies = [
 algs = [
   "optax",
   "distrax",
-  "safetensors",
   "flashbax==0.1.0",
   "wandb",
   "hydra-core>=1.3.2",

diff --git a/tests/hanabi/test_hanabi.py b/tests/hanabi/test_hanabi.py
@@ -4,7 +4,6 @@
 import jax
 from jax import numpy as jnp
 from jaxmarl import make
-from jaxmarl.wrappers.baselines import LogWrapper
 
 env = make("hanabi")
 dir_path = os.path.dirname(os.path.realpath(__file__))