accomodating substrafl api chhange

owkin · Jan 23, 2024 · d773739 · d773739
1 parent 390a466
commit d773739
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/fedeca/algorithms/torch_dp_fed_avg_algo.py b/fedeca/algorithms/torch_dp_fed_avg_algo.py
@@ -390,6 +390,7 @@ def _update_from_checkpoint(self, checkpoint: dict) -> None:
 
         # For some reason substrafl save and load client before calling train
         if "privacy_accountant_state_dict" in checkpoint:
+
             self.accountant = RDPAccountant()
             self.accountant.load_state_dict(
                 checkpoint.pop("privacy_accountant_state_dict")
@@ -427,9 +428,9 @@ def _update_from_checkpoint(self, checkpoint: dict) -> None:
         self._index_generator = checkpoint.pop("index_generator")
 
         if self._device == torch.device("cpu"):
-            torch.set_rng_state(checkpoint.pop("rng_state").to(self._device))
+            torch.set_rng_state(checkpoint.pop("torch_rng_state").to(self._device))
         else:
-            torch.cuda.set_rng_state(checkpoint.pop("rng_state").to("cpu"))
+            torch.cuda.set_rng_state(checkpoint.pop("torch_rng_state").to("cpu"))
 
         attr_names = [
             "dp_max_grad_norm",