fix: Environment.wraps signature now has correct arguments

epignatelli · Jan 13, 2024 · 88fc99c · 88fc99c
1 parent 0f07f3d
commit 88fc99c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/helx/envs/environment.py b/helx/envs/environment.py
@@ -57,8 +57,8 @@ def step(
 class EnvironmentWrapper(Environment):
     env: Any = struct.field(pytree_node=False)
 
-    @abc.abstractmethod
-    def wraps(self, env: Any) -> Timestep:
+    @abc.abstractclassmethod
+    def wraps(cls, env: Any) -> EnvironmentWrapper:
         raise NotImplementedError()
 
     def unwrapped(self) -> Any: