From b3490864e33f912589194aada30ca95ce1c882ff Mon Sep 17 00:00:00 2001 From: vmoens Date: Tue, 20 Feb 2024 17:28:56 -0800 Subject: [PATCH] fix --- test/test_rb.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/test_rb.py b/test/test_rb.py index e47d9ed58cf..a2d7988ff48 100644 --- a/test/test_rb.py +++ b/test/test_rb.py @@ -737,11 +737,11 @@ class TC: if data_type in ("tensor", "pytree"): tree_map( torch.testing.assert_close, - tree_flatten(storage._storage)[0], - tree_flatten(storage_recover._storage)[0], + tree_flatten(storage[:])[0], + tree_flatten(storage_recover[:])[0], ) else: - assert_allclose_td(storage._storage, storage_recover._storage) + assert_allclose_td(storage[:], storage_recover[:]) if data == "tc": assert storage._storage.text == storage_recover._storage.text