Fix for Resuming Training gives CheckpointMismatchError

src/pykeen/trackers/file.py

@@ -148,6 +148,13 @@ class JSONResultTracker(FileResultTracker):
     extension = "jsonl"
 
     def _write(self, obj) -> None:
+        obj = obj.copy()
+        for key, value in obj.items():
+            # Check if value is JSON serializable
+            try:
+                json.dumps(value)
+            except TypeError:
+                obj[key] = value.__class__.__name__
         print(json.dumps(obj), file=self.file, flush=True)  # noqa:T201
 
     # docstr-coverage: inherited

src/pykeen/trackers/wandb.py

@@ -49,6 +49,9 @@ def __init__(
             raise ValueError("Weights & Biases requires a project name.")
         self.project = project
 
+        if "allow_val_change" not in kwargs:
+            kwargs["allow_val_change"] = None
+
         if offline:
             os.environ[self.wandb.env.MODE] = "dryrun"  # type: ignore
         self.kwargs = kwargs
@@ -80,4 +83,8 @@ def log_params(self, params: Mapping[str, Any], prefix: Optional[str] = None) ->
         if self.run is None:
             raise AssertionError("start_run must be called before logging any metrics")
         params = flatten_dictionary(dictionary=params, prefix=prefix)
-        self.run.config.update(params)
+
+        if self.kwargs["allow_val_change"]:
+            self.run.config.update(params, allow_val_change=True)
+        else:
+            self.run.config.update(params)

src/pykeen/training/training_loop.py

@@ -1238,11 +1238,7 @@ def _load_state(
 
         logger.info(f"=> loading checkpoint '{path}'")
         checkpoint = torch.load(path)
-        if checkpoint["checksum"] != self.checksum:
-            raise CheckpointMismatchError(
-                f"The checkpoint file '{path}' that was provided already exists, but seems to be "
-                f"from a different training loop setup.",
-            )
+
         # Cuda requires its own random state, which can only be set when a cuda device is available
         torch_cuda_random_state = checkpoint["torch_cuda_random_state"]
         if torch_cuda_random_state is not None and torch.cuda.is_available():
@@ -1306,5 +1302,9 @@ def _load_state(
         np.random.set_state(checkpoint["np_random_state"])
         torch.random.set_rng_state(checkpoint["torch_random_state"])
         logger.info(f"=> loaded checkpoint '{path}' stopped after having finished epoch {checkpoint['epoch']}")
-
+        if checkpoint["checksum"] != self.checksum:
+            raise CheckpointMismatchError(
+                f"The checkpoint file '{path}' that was provided already exists, but seems to be "
+                f"from a different training loop setup.",
+            )
         return best_epoch_model_file_path, best_epoch