Make sure remote checkpoints are also saved to the current directory

src/metatrain/utils/io.py

@@ -1,3 +1,4 @@
+import shutil
 import warnings
 from pathlib import Path
 from typing import Any, Optional, Union
@@ -128,9 +129,13 @@ def load_model(
             filename += split_path[i] + "/"
         filename = filename[:-1]
         path = hf_hub_download(repo_id, filename, token=kwargs["huggingface_api_token"])
+        # make sure to copy the checkpoint to the current directory
+        shutil.copy(path, Path.cwd() / str(path).split("/")[-1])
 
     elif urlparse(str(path)).scheme:
         path, _ = urlretrieve(str(path))
+        # make sure to copy the checkpoint to the current directory
+        shutil.copy(path, Path.cwd() / str(path).split("/")[-1])
 
     else:
         pass

tests/utils/test_io.py

@@ -47,6 +47,9 @@ def test_is_exported_file():
 def test_load_model_checkpoint(path):
     model = load_model(path, architecture_name="experimental.soap_bpnn")
     assert type(model) is SoapBpnn
+    if str(path).startswith("file:"):
+        # test that the checkpoint is also copied to the current directory
+        assert Path("model-32-bit.ckpt").exists()
 
 
 @pytest.mark.parametrize(