add docstring

GFNOrg · Dec 6, 2024 · 2672d9b · 2672d9b
1 parent 958139f
commit 2672d9b
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/src/gfn/env.py b/src/gfn/env.py
@@ -94,6 +94,8 @@ def states_from_batch_shape(
 
         Args:
             batch_shape: Tuple representing the shape of the batch of states.
+            random (optional): Initalize states randomly.
+            sink (optional): States initialized with s_f (the sink state).
 
         Returns:
             States: A batch of initial states.