revert back

Author: nerdai

fl4health/clients/basic_client.py

@@ -929,9 +929,9 @@ def setup_client(self, config: Config) -> None:
         # batch_size * num_validation_steps
         self.num_val_samples = len(self.val_loader.dataset)  # type: ignore
         if self.num_validation_steps is not None:
-            assert (
-                self.val_loader.batch_size is not None
-            ), "Validation batch size must be defined if we want to limit the number of validation steps"
+            assert self.val_loader.batch_size is not None, (
+                "Validation batch size must be defined if we want to limit the number of validation steps"
+            )
             self.num_val_samples = self.num_validation_steps * self.val_loader.batch_size
 
         if self.test_loader:

fl4health/mixins/personalized/ditto.py

@@ -2,7 +2,7 @@
 
 import copy
 import warnings
-from logging import INFO, WARN
+from logging import ERROR, INFO, WARN
 from typing import Any, Protocol, cast, runtime_checkable
 
 import torch
@@ -364,7 +364,6 @@ def train_step(
             model optimization steps. The prediction dictionary contains predictions indexed a "global" and "local"
             corresponding to predictions from the global and local Ditto models for metric evaluations.
         """
-
         # global
         global_losses, global_preds = self._compute_preds_and_losses(
             self.safe_global_model(), self.optimizers["global"], input, target
@@ -406,7 +405,6 @@ def train_step(
     def val_step(
         self: DittoPersonalizedProtocol, input: TorchInputType, target: TorchTargetType
     ) -> tuple[EvaluationLosses, TorchPredType]:
-
         # global
         global_losses, global_preds = self._val_step_with_model(self.safe_global_model(), input, target)
         # local

tests/clients/test_basic_client.py

@@ -248,5 +248,4 @@ def mock_validate_or_test(  # type: ignore
             fold_loss_dict_into_metrics(self.mock_metrics, self.mock_loss_dict, logging_mode)
         if logging_mode == LoggingMode.VALIDATION:
             return self.mock_loss, self.mock_metrics
-        else:
-            return self.mock_loss, self.mock_metrics_test
+        return self.mock_loss, self.mock_metrics_test