Fixes.

Author: NikoOinonen

mlspm/utils.py

@@ -258,9 +258,9 @@ def load_checkpoint(
     import torch
 
     if rank is None:
-        state = torch.load(file_name)
+        state = torch.load(file_name, weights_only=False)
     else:
-        state = torch.load(file_name, map_location={"cuda:0": f"cuda:{rank}"})
+        state = torch.load(file_name, map_location={"cuda:0": f"cuda:{rank}"}, weights_only=False)
     model.load_state_dict(state["model_params"])
 
     if optimizer:

tests/test_data_generation.py

@@ -38,7 +38,7 @@ def test_tar_writer():
 
     with pytest.raises(RuntimeError):
         # Cannot overwrite an existing file
-        with TarWriter(base_path, base_name, max_count=10) as tar_writer:
+        with TarWriter(base_path, base_name, max_count=10, async_write=False) as tar_writer:
             pass
 
     rmtree(base_path)