removed .jax module prefixes from imports in pymdp.envs.rollout

infer-actively · Oct 2, 2024 · 99cf06e · 99cf06e
1 parent b4513b4
commit 99cf06e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pymdp/envs/rollout.py b/pymdp/envs/rollout.py
@@ -3,8 +3,8 @@
 import jax.tree_util as jtu
 import jax.lax
 
-from pymdp.jax.agent import Agent
-from pymdp.jax.envs.env import Env
+from pymdp.agent import Agent
+from pymdp.envs.env import Env
 
 
 def rollout(agent: Agent, env: Env, num_timesteps: int, rng_key: jr.PRNGKey, policy_search=None):