fix

pytorch · Feb 21, 2024 · b349086 · b349086
1 parent fb9e562
commit b349086
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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