chore: use asserts instead of ignores

instadeepai · Jul 23, 2024 · 6579542 · 6579542
1 parent 766fab7
commit 6579542
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/flashbax/buffers/prioritised_trajectory_buffer.py b/flashbax/buffers/prioritised_trajectory_buffer.py
@@ -798,10 +798,11 @@ 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,
-        max_length_time_axis=max_length_time_axis,  # type: ignore
+        max_length_time_axis=max_length_time_axis,
         period=period,
     )
     add_fn = functools.partial(

diff --git a/flashbax/buffers/trajectory_buffer.py b/flashbax/buffers/trajectory_buffer.py
@@ -585,10 +585,11 @@ 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,
-        max_length_time_axis=max_length_time_axis,  # type: ignore
+        max_length_time_axis=max_length_time_axis,
     )
     add_fn = functools.partial(
         add,