refactor: scaffold compatibility and docstrings

substrafl/algorithms/pytorch/torch_fed_pca_algo.py

@@ -345,12 +345,12 @@ def _update_from_checkpoint(self, checkpoint: dict) -> None:
         to be popped.
 
         Args:
-            path (pathlib.Path): path where the checkpoint is saved
+            checkpoint (dict): the checkpoint to load.
 
         Returns:
             dict: checkpoint
         """
-        super()._update_from_checkpoint(checkpoint)
+        super()._update_from_checkpoint(checkpoint=checkpoint)
         self.local_mean = checkpoint.pop("mean")
         self.local_covmat = checkpoint.pop("covariance_matrix")
         return

substrafl/algorithms/pytorch/torch_scaffold_algo.py

@@ -1,6 +1,5 @@
 import logging
 from enum import IntEnum
-from pathlib import Path
 from typing import Any
 from typing import List
 from typing import Optional
@@ -493,18 +492,18 @@ def _get_state_to_save(self) -> dict:
         )
         return local_state
 
-    def _update_from_checkpoint(self, path: Path) -> dict:
+    def _update_from_checkpoint(self, checkpoint: dict) -> None:
         """Load the local state from the checkpoint.
 
         Args:
-            path (pathlib.Path): path where the checkpoint is saved
+            checkpoint (dict): the checkpoint to load.
 
         Returns:
             dict: checkpoint
         """
-        checkpoint = super()._update_from_checkpoint(path=path)
+        checkpoint = super()._update_from_checkpoint(checkpoint=checkpoint)
         self._client_control_variate = checkpoint.pop("client_control_variate")
-        return checkpoint
+        return
 
     def summary(self):
         """Summary of the class to be exposed in the experiment summary file