Merge branch 'chore/mypy-max-len-time-axis' into chore/new-gh-runners

instadeepai · Jul 26, 2024 · 24c666b · 24c666b
2 parents 28a482c + 0ee266f
commit 24c666b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/flashbax/buffers/prioritised_trajectory_buffer.py b/flashbax/buffers/prioritised_trajectory_buffer.py
@@ -798,6 +798,7 @@ def make_prioritised_trajectory_buffer(
     if max_size is not None:
         max_length_time_axis = max_size // add_batch_size
 
+    assert max_length_time_axis is not None
     init_fn = functools.partial(
         prioritised_init,
         add_batch_size=add_batch_size,

diff --git a/flashbax/buffers/trajectory_buffer.py b/flashbax/buffers/trajectory_buffer.py
@@ -585,6 +585,7 @@ def make_trajectory_buffer(
     if max_size is not None:
         max_length_time_axis = max_size // add_batch_size
 
+    assert max_length_time_axis is not None
     init_fn = functools.partial(
         init,
         add_batch_size=add_batch_size,