Make mypy happy and use [] instea of . notation

* mypy does not know that value of keys are Batch
dantp-ai · Mar 18, 2024 · a481e2f · a481e2f
1 parent 91a3ee2
commit a481e2f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/base/test_buffer.py b/test/base/test_buffer.py
@@ -190,9 +190,9 @@ def test_ignore_obs_next(size: int = 10) -> None:
                 ],
             ),
         )
-    assert np.allclose(data.info["if"], data2.info["if"])
+    assert np.allclose(data["info"]["if"], data2["info"]["if"])
     assert np.allclose(
-        data.info["if"],
+        data["info"]["if"],
         np.array(
             [
                 [0, 0, 0, 0],