Update src/sbv2/utils.py

src/sbv2/utils.py

@@ -1,12 +1,7 @@
-import torch
-
-def load_checkpoint(filepath, model, optimizer=None, strict=True):
-    print(f"Loading checkpoint: {filepath}")
-    checkpoint_dict = torch.load(filepath, map_location="cpu")
-    model.load_state_dict(checkpoint_dict["model"], strict=strict)
-
-    if optimizer is not None and "optimizer" in checkpoint_dict:
-        optimizer.load_state_dict(checkpoint_dict["optimizer"])
-
-    print("Checkpoint loaded.")
-    return model
+from safetensors.torch import load_file
+
+def load_checkpoint(filepath, model, strict=True):
+    print(f"Loading checkpoint: {filepath}")
+    checkpoint_dict = load_file(filepath)
+    model.load_state_dict(checkpoint_dict, strict=strict)
+    return model