@@ -672,6 +672,8 @@ def test_storage_state_dict(self, storage_in, storage_out, init_out, backend):
672
672
def test_storage_dumps_loads (
673
673
self , device_data , storage_type , data_type , isinit , tmpdir
674
674
):
675
+ torch .manual_seed (0 )
676
+
675
677
dir_rb = tmpdir / "rb"
676
678
dir_save = tmpdir / "save"
677
679
dir_rb .mkdir ()
@@ -716,25 +718,30 @@ class TC:
716
718
)
717
719
else :
718
720
raise NotImplementedError
721
+
719
722
if storage_type in (LazyMemmapStorage ,):
720
723
storage = storage_type (max_size = 10 , scratch_dir = dir_rb )
721
724
else :
722
725
storage = storage_type (max_size = 10 )
726
+
723
727
# We cast the device to CPU as CUDA isn't automatically cast to CPU when using range() index
724
728
if data_type == "pytree" :
725
729
storage .set (range (3 ), tree_map (lambda x : x .cpu (), data ))
726
730
else :
727
731
storage .set (range (3 ), data .cpu ())
732
+
728
733
storage .dumps (dir_save )
729
734
# check we can dump twice
730
735
storage .dumps (dir_save )
731
736
732
737
storage_recover = storage_type (max_size = 10 )
733
738
if isinit :
734
739
if data_type == "pytree" :
735
- storage_recover .set (range (3 ), tree_map (lambda x : x .cpu ().zero_ (), data ))
740
+ storage_recover .set (
741
+ range (3 ), tree_map (lambda x : x .cpu ().clone ().zero_ (), data )
742
+ )
736
743
else :
737
- storage_recover .set (range (3 ), data .cpu ().zero_ ())
744
+ storage_recover .set (range (3 ), data .cpu ().clone (). zero_ ())
738
745
739
746
if data_type in ("tensor" , "pytree" ) and not isinit :
740
747
with pytest .raises (
0 commit comments