diff --git a/examples/nnunet_pfl_example/README.md b/examples/nnunet_pfl_example/README.md
new file mode 100644
index 000000000..159b21beb
--- /dev/null
+++ b/examples/nnunet_pfl_example/README.md
@@ -0,0 +1,19 @@
+# NnUNetClient With Personalization Example
+
+Building on the [nnunet_example](../nnunet_example/README.md), here we demonstrate how personalized
+methods can be applied to `NnUNetClient` class. The config requirements remain the same as in the original
+`nnunet_example`.
+
+To run a federated learning experiment with nnunet models, first ensure you are in the FL4Health directory and then start the nnunet server using the following command. To view a list of optional flags use the --help flag
+
+```bash
+python -m examples.nnunet_pfl_example.server --config_path examples/nnunet_pfl_example/config.yaml
+```
+
+Once the server has started, start the necessary number of clients specified by the n_clients key in the config file. Each client can be started by running the following command in a separate session. To view a list of optional flags use the --help flag.
+
+```bash
+python -m examples.nnunet_pfl_example.client --dataset_path examples/datasets/nnunet --personalized-strategy ditto
+```
+
+The same MSD dataset that was used in the original `nnunet_example` is also used here.
diff --git a/examples/nnunet_pfl_example/__init__.py b/examples/nnunet_pfl_example/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/examples/nnunet_pfl_example/client.py b/examples/nnunet_pfl_example/client.py
new file mode 100644
index 000000000..eab6ebfe3
--- /dev/null
+++ b/examples/nnunet_pfl_example/client.py
@@ -0,0 +1,245 @@
+import argparse
+import os
+import warnings
+from logging import DEBUG, INFO
+from os.path import exists, join
+from pathlib import Path
+from typing import Any, Literal
+
+from fl4health.checkpointing.checkpointer import PerRoundStateCheckpointer
+from fl4health.checkpointing.client_module import ClientCheckpointAndStateModule
+
+with warnings.catch_warnings():
+    # Silence deprecation warnings from sentry sdk due to flwr and wandb
+    # https://github.com/adap/flower/issues/4086
+    warnings.filterwarnings("ignore", category=DeprecationWarning)
+    import wandb  # noqa: F401
+
+import torch
+from flwr.client import start_client
+from flwr.common.logger import log, update_console_handler
+from nnunetv2.dataset_conversion.convert_MSD_dataset import convert_msd_dataset
+from torchmetrics.segmentation import GeneralizedDiceScore
+
+from fl4health.clients.nnunet_client import NnunetClient
+from fl4health.metrics import TorchMetric
+from fl4health.metrics.compound_metrics import TransformsMetric
+from fl4health.mixins.personalized import PersonalizedMode, make_it_personal
+from fl4health.utils.load_data import load_msd_dataset
+from fl4health.utils.msd_dataset_sources import get_msd_dataset_enum, msd_num_labels
+from fl4health.utils.nnunet_utils import get_segs_from_probs, set_nnunet_env
+
+personalized_client_classes = {"ditto": make_it_personal(NnunetClient, PersonalizedMode.DITTO)}
+
+
+def main(
+    dataset_path: Path,
+    msd_dataset_name: str,
+    server_address: str,
+    fold: int | str,
+    always_preprocess: bool = False,
+    verbose: bool = True,
+    compile: bool = True,
+    intermediate_client_state_dir: str | None = None,
+    client_name: str | None = None,
+    personalized_strategy: Literal["ditto"] | None = None,
+) -> None:
+    with torch.autograd.set_detect_anomaly(True):
+        # Log device and server address
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        log(INFO, f"Using device: {device}")
+        log(INFO, f"Using server address: {server_address}")
+
+        # Load the dataset if necessary
+        msd_dataset_enum = get_msd_dataset_enum(msd_dataset_name)
+        nnUNet_raw = join(dataset_path, "nnunet_raw")
+        if not exists(join(nnUNet_raw, msd_dataset_enum.value)):
+            log(INFO, f"Downloading and extracting {msd_dataset_enum.value} dataset")
+            load_msd_dataset(nnUNet_raw, msd_dataset_name)
+
+        # The dataset ID will be the same as the MSD Task number
+        dataset_id = int(msd_dataset_enum.value[4:6])
+        nnunet_dataset_name = f"Dataset{dataset_id:03d}_{msd_dataset_enum.value.split('_')[1]}"
+
+        # Convert the msd dataset if necessary
+        if not exists(join(nnUNet_raw, nnunet_dataset_name)):
+            log(INFO, f"Converting {msd_dataset_enum.value} into nnunet dataset")
+            convert_msd_dataset(source_folder=join(nnUNet_raw, msd_dataset_enum.value))
+
+        # Create a metric
+        dice = TransformsMetric(
+            metric=TorchMetric(
+                name="Pseudo DICE",
+                metric=GeneralizedDiceScore(
+                    num_classes=msd_num_labels[msd_dataset_enum], weight_type="square", include_background=False
+                ).to(device),
+            ),
+            pred_transforms=[torch.sigmoid, get_segs_from_probs],
+        )
+
+        if intermediate_client_state_dir is not None:
+            checkpoint_and_state_module = ClientCheckpointAndStateModule(
+                state_checkpointer=PerRoundStateCheckpointer(Path(intermediate_client_state_dir))
+            )
+        else:
+            checkpoint_and_state_module = None
+
+        # Create client
+        client_kwargs: dict[str, Any] = {}
+        client_kwargs.update(
+            # Args specific to nnUNetClient
+            dataset_id=dataset_id,
+            fold=fold,
+            always_preprocess=always_preprocess,
+            verbose=verbose,
+            compile=compile,
+            # BaseClient Args
+            device=device,
+            metrics=[dice],
+            progress_bar=verbose,
+            checkpoint_and_state_module=checkpoint_and_state_module,
+            client_name=client_name,
+        )
+        if personalized_strategy:
+            log(INFO, f"Setting up client for personalized strategy: {personalized_strategy}")
+            client = personalized_client_classes[personalized_strategy](**client_kwargs)
+        else:
+            log(INFO, "Setting up client without personalization")
+            client = NnunetClient(**client_kwargs)
+        log(INFO, f"Using client: {type(client).__name__}")
+
+        start_client(server_address=server_address, client=client.to_client())
+
+        # Shutdown the client
+        client.shutdown()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        prog="nnunet_example/client.py",
+        description="An exampled of nnUNetClient on any of the Medical \
+            Segmentation Decathlon (MSD) datasets. Automatically generates a \
+            nnunet segmentation model and trains it in a federated setting",
+    )
+
+    # I have to use underscores instead of dashes because thats how they
+    # defined it in run_smoke_tests
+    parser.add_argument(
+        "--dataset_path",
+        type=str,
+        required=True,
+        help="Path to the folder in which data should be stored. This script \
+            will automatically create nnunet_raw, and nnunet_preprocessed \
+            subfolders if they don't already exist. This script will also \
+            attempt to download and prepare the MSD Dataset into the \
+            nnunet_raw folder if it does not already exist.",
+    )
+    parser.add_argument(
+        "--fold",
+        type=str,
+        required=False,
+        default="0",
+        help="[OPTIONAL] Which fold of the local client dataset to use for \
+            validation. nnunet defaults to 5 folds (0 to 4). Can also be set \
+            to 'all' to use all the data for both training and validation. \
+            Defaults to fold 0",
+    )
+    parser.add_argument(
+        "--msd_dataset_name",
+        type=str,
+        required=False,
+        default="Task04_Hippocampus",  # The smallest dataset
+        help="[OPTIONAL] Name of the MSD dataset to use. The options are \
+            defined by the values of the MsdDataset enum as returned by the \
+            get_msd_dataset_enum function",
+    )
+    parser.add_argument(
+        "--always-preprocess",
+        action="store_true",
+        required=False,
+        help="[OPTIONAL] Use this to force preprocessing the nnunet data \
+            even if the preprocessed data is found to already exist",
+    )
+    parser.add_argument(
+        "--server_address",
+        type=str,
+        required=False,
+        default="0.0.0.0:8080",
+        help="[OPTIONAL] The server address for the clients to communicate \
+            to the server through. Defaults to 0.0.0.0:8080",
+    )
+    parser.add_argument(
+        "--verbose",
+        action="store_true",
+        required=False,
+        help="[OPTIONAL] Use this flag to see extra INFO logs and a progress bar",
+    )
+    parser.add_argument(
+        "--debug",
+        help="[OPTIONAL] Include flag to print DEBUG logs",
+        action="store_const",
+        dest="logLevel",
+        const=DEBUG,
+        default=INFO,
+    )
+    parser.add_argument(
+        "--skip-compile",
+        action="store_true",
+        required=False,
+        help="[OPTIONAL] Include flag to train without jit compiling the pytorch model first",
+    )
+
+    parser.add_argument(
+        "--intermediate-client-state-dir",
+        type=str,
+        required=False,
+        default=None,
+        help="[OPTIONAL] Directory to store client state during training. Defaults to None",
+    )
+    parser.add_argument(
+        "--client-name",
+        type=str,
+        required=False,
+        default=None,
+        help="[OPTIONAL] Name of the client used to name client state checkpoint. \
+        Defaults to None, in which case a random name is generated for the client",
+    )
+    parser.add_argument(
+        "--personalized-strategy",
+        type=str,
+        required=False,
+        default=None,
+        help="[OPTIONAL] Personalized strategy to use. For now, can only be 'ditto'. \
+        Defaults to None, in which no personalized strategy is applied.",
+    )
+
+    args = parser.parse_args()
+
+    # Set the log level
+    update_console_handler(level=args.logLevel)
+
+    # Create nnunet directory structure and set environment variables
+    nnUNet_raw = join(args.dataset_path, "nnunet_raw")
+    nnUNet_preprocessed = join(args.dataset_path, "nnunet_preprocessed")
+    os.makedirs(nnUNet_raw, exist_ok=True)
+    os.makedirs(nnUNet_preprocessed, exist_ok=True)
+    set_nnunet_env(
+        nnUNet_raw=nnUNet_raw,
+        nnUNet_preprocessed=nnUNet_preprocessed,
+        nnUNet_results=join(args.dataset_path, "nnUNet_results"),
+    )
+
+    # Check fold argument and start main method
+    fold: int | str = "all" if args.fold == "all" else int(args.fold)
+    main(
+        dataset_path=Path(args.dataset_path),
+        msd_dataset_name=args.msd_dataset_name,
+        server_address=args.server_address,
+        fold=fold,
+        always_preprocess=args.always_preprocess,
+        verbose=args.verbose,
+        compile=not args.skip_compile,
+        intermediate_client_state_dir=args.intermediate_client_state_dir,
+        client_name=args.client_name,
+        personalized_strategy=args.personalized_strategy,
+    )
diff --git a/examples/nnunet_pfl_example/config.yaml b/examples/nnunet_pfl_example/config.yaml
new file mode 100644
index 000000000..d2371d0d6
--- /dev/null
+++ b/examples/nnunet_pfl_example/config.yaml
@@ -0,0 +1,8 @@
+# Parameters that describe the server
+n_server_rounds: 1
+
+# Parameters that describe the clients
+n_clients: 1
+local_epochs: 1
+
+nnunet_config: 2d
diff --git a/examples/nnunet_pfl_example/server.py b/examples/nnunet_pfl_example/server.py
new file mode 100644
index 000000000..1cfe9fa8f
--- /dev/null
+++ b/examples/nnunet_pfl_example/server.py
@@ -0,0 +1,184 @@
+import argparse
+import json
+import pickle
+from functools import partial
+from pathlib import Path
+
+import flwr as fl
+import torch
+import yaml
+from flwr.common.parameter import ndarrays_to_parameters
+from flwr.common.typing import Config
+from flwr.server.client_manager import SimpleClientManager
+
+from fl4health.checkpointing.checkpointer import PerRoundStateCheckpointer
+from fl4health.checkpointing.server_module import NnUnetServerCheckpointAndStateModule
+from fl4health.metrics.metric_aggregation import evaluate_metrics_aggregation_fn, fit_metrics_aggregation_fn
+from fl4health.parameter_exchange.full_exchanger import FullParameterExchanger
+from fl4health.parameter_exchange.packing_exchanger import FullParameterExchangerWithPacking
+from fl4health.parameter_exchange.parameter_packer import (
+    ParameterPackerAdaptiveConstraint,
+)
+from fl4health.servers.nnunet_server import NnunetServer
+from fl4health.strategies.fedavg_with_adaptive_constraint import FedAvgWithAdaptiveConstraint
+from fl4health.utils.config import make_dict_with_epochs_or_steps
+from fl4health.utils.random import set_all_random_seeds
+
+
+def get_config(
+    current_server_round: int,
+    nnunet_config: str,
+    n_server_rounds: int,
+    batch_size: int,
+    n_clients: int,
+    nnunet_plans: str | None = None,
+    local_epochs: int | None = None,
+    local_steps: int | None = None,
+) -> Config:
+    # Create config
+    config: Config = {
+        "n_clients": n_clients,
+        "nnunet_config": nnunet_config,
+        "n_server_rounds": n_server_rounds,
+        "batch_size": batch_size,
+        **make_dict_with_epochs_or_steps(local_epochs, local_steps),
+        "current_server_round": current_server_round,
+    }
+
+    # Check if plans were provided
+    if nnunet_plans is not None:
+        plans_bytes = pickle.dumps(json.load(open(nnunet_plans, "r")))
+        config["nnunet_plans"] = plans_bytes
+
+    return config
+
+
+def main(
+    config: dict,
+    server_address: str,
+    intermediate_server_state_dir: str | None = None,
+    server_name: str | None = None,
+) -> None:
+    # Partial function with everything set except current server round
+    fit_config_fn = partial(
+        get_config,
+        n_clients=config["n_clients"],
+        nnunet_config=config["nnunet_config"],
+        n_server_rounds=config["n_server_rounds"],
+        batch_size=0,  # Set this to 0 because we're not using it
+        nnunet_plans=config.get("nnunet_plans"),
+        local_epochs=config.get("local_epochs"),
+        local_steps=config.get("local_steps"),
+    )
+
+    if config.get("starting_checkpoint"):
+        model = torch.load(config["starting_checkpoint"])
+        # Of course nnunet stores their pytorch models differently.
+        params = ndarrays_to_parameters([val.cpu().numpy() for _, val in model["network_weights"].items()])
+    else:
+        params = None
+
+    strategy = FedAvgWithAdaptiveConstraint(
+        min_fit_clients=config["n_clients"],
+        min_evaluate_clients=config["n_clients"],
+        # Server waits for min_available_clients before starting FL rounds
+        min_available_clients=config["n_clients"],
+        on_fit_config_fn=fit_config_fn,
+        # We use the same fit config function, as nothing changes for eval
+        on_evaluate_config_fn=fit_config_fn,
+        fit_metrics_aggregation_fn=fit_metrics_aggregation_fn,
+        evaluate_metrics_aggregation_fn=evaluate_metrics_aggregation_fn,
+        initial_parameters=params,
+        initial_loss_weight=0.1,
+        adapt_loss_weight=False,
+    )
+
+    state_checkpointer = (
+        PerRoundStateCheckpointer(Path(intermediate_server_state_dir))
+        if intermediate_server_state_dir is not None
+        else None
+    )
+    checkpoint_and_state_module = NnUnetServerCheckpointAndStateModule(
+        model=None, parameter_exchanger=FullParameterExchanger(), state_checkpointer=state_checkpointer
+    )
+
+    checkpoint_and_state_module.parameter_exchanger = FullParameterExchangerWithPacking(  # type:ignore [assignment]
+        ParameterPackerAdaptiveConstraint()
+    )
+
+    server = NnunetServer(
+        client_manager=SimpleClientManager(),
+        fl_config=config,
+        # The fit_config_fn contains all of the necessary information for param initialization, so we reuse it here
+        on_init_parameters_config_fn=fit_config_fn,
+        strategy=strategy,
+        checkpoint_and_state_module=checkpoint_and_state_module,
+        server_name=server_name,
+        accept_failures=False,
+    )
+
+    fl.server.start_server(
+        server=server,
+        server_address=server_address,
+        config=fl.server.ServerConfig(num_rounds=config["n_server_rounds"]),
+    )
+
+    # Shutdown server
+    server.shutdown()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--config_path",
+        action="store",
+        type=str,
+        help="Path to the configuration file. See examples/nnunet_example/README.md for more info",
+    )
+    parser.add_argument(
+        "--server-address",
+        type=str,
+        required=False,
+        default="0.0.0.0:8080",
+        help="[OPTIONAL] The address to use for the server. Defaults to \
+        0.0.0.0:8080",
+    )
+
+    parser.add_argument(
+        "--intermediate-server-state-dir",
+        type=str,
+        required=False,
+        default=None,
+        help="[OPTIONAL] Directory to store server state. Defaults to None",
+    )
+    parser.add_argument(
+        "--server-name",
+        type=str,
+        required=False,
+        default=None,
+        help="[OPTIONAL] Name of the server used as file name when \
+            checkpointing server state. Defaults to \
+            None, in which case the server will generate random name \
+            ",
+    )
+    parser.add_argument(
+        "--seed",
+        action="store",
+        type=int,
+        help="Seed for the random number generators across python, torch, and numpy",
+        required=False,
+    )
+    args = parser.parse_args()
+
+    with open(args.config_path, "r") as f:
+        config = yaml.safe_load(f)
+
+    # Set the random seed for reproducibility
+    set_all_random_seeds(args.seed)
+
+    main(
+        config,
+        server_address=args.server_address,
+        intermediate_server_state_dir=args.intermediate_server_state_dir,
+        server_name=args.server_name,
+    )
diff --git a/fl4health/clients/basic_client.py b/fl4health/clients/basic_client.py
index eaff0df1d..258a7c2a8 100644
--- a/fl4health/clients/basic_client.py
+++ b/fl4health/clients/basic_client.py
@@ -559,38 +559,103 @@ def update_metric_manager(
         """
         metric_manager.update(preds, target)
 
-    def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]:
-        """
-        Given a single batch of input and target data, generate predictions, compute loss, update parameters and
-        optionally update metrics if they exist. (i.e. backprop on a single batch of data).
-        Assumes ``self.model`` is in train mode already.
+    def _compute_preds_and_losses(
+        self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        """Helper method within the train step for computing preds and losses.
+
+        NOTE: Subclasses should implement this helper method if there is a need
+        to specialize this part of the overall train step.
 
         Args:
+            model (nn.Module): the model used to make predictions
+            optimizer (Optimizer): the associated optimizer
             input (TorchInputType): The input to be fed into the model.
             target (TorchTargetType): The target corresponding to the input.
 
         Returns:
             tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with
-            a dictionary of any predictions produced by the model.
+            a dictionary of any predictions produced by the model prior to the
+            application of the backwards phase
         """
         # Clear gradients from optimizer if they exist
-        self.optimizers["global"].zero_grad()
+        optimizer.zero_grad()
 
         # Call user defined methods to get predictions and compute loss
-        preds, features = self.predict(input)
+        preds, features = self._predict_with_model(model, input)
         target = self.transform_target(target)
         losses = self.compute_training_loss(preds, features, target)
 
+        return losses, preds
+
+    def _apply_backwards_on_losses_and_take_step(
+        self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses
+    ) -> TrainingLosses:
+        """Helper method within the train step for applying backwards on losses and taking step with optimizer.
+
+        NOTE: Subclasses should implement this helper method if there is a need
+        to specialize this part of the overall train step.
+
+        Args:
+            model (nn.Module): the model used for making predictions. Passed here in case subclasses need it.
+            optimizer (Optimizer): the optimizer with which we take the step
+            losses (TrainingLosses): the losses to apply backwards on
+
+        Returns:
+            TrainingLosses: The losses object post backwards application
+        """
         # Compute backward pass and update parameters with optimizer
         losses.backward["backward"].backward()
         self.transform_gradients(losses)
-        self.optimizers["global"].step()
+        optimizer.step()
+
+        return losses
+
+    def _train_step_with_model_and_optimizer(
+        self, model: torch.nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        """Helper train step method that allows for injection of model and optimizer.
+
+        NOTE: Subclasses should implement this method if there is a need to specialize
+        the train_step logic.
+
+        Args:
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with
+            a dictionary of any predictions produced by the model.
+        """
+
+        losses, preds = self._compute_preds_and_losses(model, optimizer, input, target)
+        losses = self._apply_backwards_on_losses_and_take_step(model, optimizer, losses)
 
         return losses, preds
 
-    def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]:
+    def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]:
         """
-        Given input and target, compute loss, update loss and metrics. Assumes ``self.model`` is in eval mode already.
+        Given a single batch of input and target data, generate predictions, compute loss, update parameters and
+        optionally update metrics if they exist. (i.e. backprop on a single batch of data).
+        Assumes ``self.model`` is in train mode already.
+
+        Args:
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with
+            a dictionary of any predictions produced by the model.
+        """
+        return self._train_step_with_model_and_optimizer(self.model, self.optimizers["global"], input, target)
+
+    def _val_step_with_model(
+        self, model: nn.Module, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[EvaluationLosses, TorchPredType]:
+        """Helper method for val_step that allows for injection of model.
+
+        NOTE: Subclasses should implement this method if there is a need to
+        specialize the val_step logic.
 
         Args:
             input (TorchInputType): The input to be fed into the model.
@@ -603,12 +668,27 @@ def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Eval
 
         # Get preds and compute loss
         with torch.no_grad():
-            preds, features = self.predict(input)
+            preds, features = self._predict_with_model(model, input)
             target = self.transform_target(target)
             losses = self.compute_evaluation_loss(preds, features, target)
 
         return losses, preds
 
+    def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]:
+        """
+        Given input and target, compute loss, update loss and metrics. Assumes ``self.model`` is in eval mode already.
+
+        Args:
+            input (TorchInputType): The input to be fed into the model.
+            target (TorchTargetType): The target corresponding to the input.
+
+        Returns:
+            tuple[EvaluationLosses, TorchPredType]: The losses object from the val step along with a dictionary of the
+            predictions produced by the model.
+        """
+
+        return self._val_step_with_model(self.model, input, target)
+
     def train_by_epochs(
         self,
         epochs: int,
@@ -975,14 +1055,19 @@ def get_parameter_exchanger(self, config: Config) -> ParameterExchanger:
         """
         return FullParameterExchanger()
 
-    def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
-        """
-        Computes the prediction(s), and potentially features, of the model(s) given the input.
+    def _predict_with_model(
+        self, model: torch.nn.Module, input: TorchInputType
+    ) -> tuple[TorchPredType, TorchFeatureType]:
+        """Helper predict method that allows for injection of model.
+
+        NOTE: Subclasses should implement this method if there is need to specialize
+        the predict logic of the client.
 
         Args:
+            model (torch.nn.Module): the model with which to make predictions
             input (TorchInputType): Inputs to be fed into the model. If input is of type ``dict[str, torch.Tensor]``,
                 it is assumed that the keys of input match the names of the keyword arguments of
-                ``self.model.forward().``
+                ``self.model.forward().`
 
         Returns:
             tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains a dictionary of
@@ -996,13 +1081,14 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp
             ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model
                 forward.
         """
+
         if isinstance(input, torch.Tensor):
-            output = self.model(input)
+            output = model(input)
         elif isinstance(input, dict):
             # If input is a dictionary, then we unpack it before computing the forward pass.
             # Note that this assumes the keys of the input match (exactly) the keyword args
             # of self.model.forward().
-            output = self.model(**input)
+            output = model(**input)
         else:
             raise TypeError("'input' must be of type torch.Tensor or dict[str, torch.Tensor].")
 
@@ -1018,6 +1104,29 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp
         else:
             raise ValueError("Model forward did not return a tensor, dictionary of tensors, or tuple of tensors")
 
+    def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
+        """
+        Computes the prediction(s), and potentially features, of the model(s) given the input.
+
+        Args:
+            input (TorchInputType): Inputs to be fed into the model. If input is of type ``dict[str, torch.Tensor]``,
+                it is assumed that the keys of input match the names of the keyword arguments of
+                ``self.model.forward().``
+
+        Returns:
+            tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains a dictionary of
+            predictions indexed by name and the second element contains intermediate activations indexed by name. By
+            passing features, we can compute losses such as the contrastive loss in MOON. All predictions included in
+            dictionary will by default be used to compute metrics separately.
+
+        Raises:
+            TypeError: Occurs when something other than a tensor or dict of tensors is passed in to the model's
+                forward method.
+            ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model
+                forward.
+        """
+        return self._predict_with_model(self.model, input)
+
     def compute_loss_and_additional_losses(
         self, preds: TorchPredType, features: TorchFeatureType, target: TorchTargetType
     ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]:
@@ -1280,6 +1389,19 @@ def update_before_epoch(self, epoch: int) -> None:
         """
         pass
 
+    def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None:
+        """Helper transform gradients method that allows for injection of model.
+
+        NOTE: Subclasses should implement this helper should there be a need to specialize the logic
+        for transforming gradients.
+
+        Args:
+            model (torch.nn.Module): the model used to generate predictions to compute losses
+            losses (TrainingLosses): The losses object from the train step
+        """
+
+        pass
+
     def transform_gradients(self, losses: TrainingLosses) -> None:
         """
         Hook function for model training only called after backwards pass but before optimizer step. Useful for
@@ -1288,7 +1410,7 @@ def transform_gradients(self, losses: TrainingLosses) -> None:
         Args:
             losses (TrainingLosses): The losses object from the train step
         """
-        pass
+        return self._transform_gradients_with_model(self.model, losses)
 
     def _save_client_state(self) -> None:
         """
diff --git a/fl4health/clients/ditto_client.py b/fl4health/clients/ditto_client.py
index 03749d4aa..00e4f4a5b 100644
--- a/fl4health/clients/ditto_client.py
+++ b/fl4health/clients/ditto_client.py
@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from logging import INFO
+from logging import DEBUG, INFO
 from pathlib import Path
 
 import torch
@@ -254,6 +254,7 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr
 
         # Forward pass on both the global and local models
         preds, features = self.predict(input)
+        log(DEBUG, f"PRED KEYS from train_step: {preds.keys()}")
         target = self.transform_target(target)  # Apply transformation (Defaults to identity)
 
         # Compute all relevant losses
@@ -270,6 +271,15 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr
         # Return dictionary of predictions where key is used to name respective MetricMeters
         return losses, preds
 
+    def val_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[EvaluationLosses, TorchPredType]:
+        # Get preds and compute loss
+        with torch.no_grad():
+            preds, features = self.predict(input)
+            target = self.transform_target(target)
+            losses = self.compute_evaluation_loss(preds, features, target)
+
+        return losses, preds
+
     def predict(
         self,
         input: TorchInputType,
@@ -328,6 +338,7 @@ def compute_loss_and_additional_losses(
         """
 
         # Compute global model vanilla loss
+        log(DEBUG, f"PRED KEYS from compute_loss_and_additional_losses: {preds.keys()}")
         assert "global" in preds
         global_loss = self.criterion(preds["global"], target)
 
@@ -366,6 +377,7 @@ def compute_training_loss(
         assert self.global_model.training and self.model.training
 
         # local loss is stored in loss, global model loss is stored in additional losses.
+        log(DEBUG, f"PRED KEYS: {preds.keys()}")
         loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target)
 
         # Setting the adaptation loss to that of the local model, as its performance should dictate whether more or
diff --git a/fl4health/clients/nnunet_client.py b/fl4health/clients/nnunet_client.py
index ecf9b57bf..ea474d2f3 100644
--- a/fl4health/clients/nnunet_client.py
+++ b/fl4health/clients/nnunet_client.py
@@ -205,12 +205,10 @@ def __init__(
             if self.verbose:
                 log(INFO, "Disabling model optimizations and JIT compilation. This may impact runtime performance.")
 
-    def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[TrainingLosses, TorchPredType]:
+    def _compute_preds_and_losses(
+        self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
         """
-        Given a single batch of input and target data, generate predictions, compute loss, update parameters and
-        optionally update metrics if they exist. (i.e. backprop on a single batch of data). Assumes ``self.model`` is
-        in train mode already.
-
         Overrides parent to include mixed precision training (autocasting and corresponding gradient scaling) as per
         the original ``nnUNetTrainer``.
 
@@ -222,32 +220,50 @@ def train_step(self, input: TorchInputType, target: TorchTargetType) -> tuple[Tr
             Tuple[TrainingLosses, TorchPredType]: The losses object from the train step along with a dictionary of any
             predictions produced by the model.
         """
+
         # If the device type is not cuda, we don't use mixed precision training and therefore can use parent method.
         if self.device.type != "cuda":
-            return super().train_step(input, target)
+            return super()._compute_preds_and_losses(model, optimizer, input, target)
 
         # As in the nnUNetTrainer, we implement mixed precision using torch.autocast and torch.GradScaler
         # Clear gradients from optimizer if they exist
-        self.optimizers["global"].zero_grad()
+        optimizer.zero_grad()
 
         # Call user defined methods to get predictions and compute loss
-        preds, features = self.predict(input)
+        preds, features = self._predict_with_model(model, input)
         target = self.transform_target(target)
         losses = self.compute_training_loss(preds, features, target)
 
+        return losses, preds
+
+    def _apply_backwards_on_losses_and_take_step(
+        self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses
+    ) -> TrainingLosses:
+        """
+        Overrides parent to make use grad scaler with ``nnUNetTrainer``.
+
+        Args:
+            model (nn.Module): the model used for making predictions. Passed here in case subclasses need it.
+            optimizer (Optimizer): the optimizer with which we take the step
+            losses (TrainingLosses): the losses to apply backwards on
+
+        Returns:
+            TrainingLosses: The losses object after backwards application
+        """
+
         # Compute scaled loss and perform backward pass
         scaled_backward_loss = self.grad_scaler.scale(losses.backward["backward"])
         scaled_backward_loss.backward()
 
         # Rescale gradients then clip based on specified norm
-        self.grad_scaler.unscale_(self.optimizers["global"])
-        self.transform_gradients(losses)
+        self.grad_scaler.unscale_(optimizer)
+        self._transform_gradients_with_model(model, losses)
 
         # Update parameters and scaler
-        self.grad_scaler.step(self.optimizers["global"])
+        self.grad_scaler.step(optimizer)
         self.grad_scaler.update()
 
-        return losses, preds
+        return losses
 
     @use_default_signal_handlers  # Dataloaders use multiprocessing
     def get_data_loaders(self, config: Config) -> tuple[DataLoader, DataLoader]:
@@ -606,7 +622,9 @@ def setup_client(self, config: Config) -> None:
         # We have to call parent method after setting up nnunet trainer
         super().setup_client(config)
 
-    def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch.Tensor]]:
+    def _predict_with_model(
+        self, model: torch.nn.Module, input: TorchInputType
+    ) -> tuple[TorchPredType, dict[str, torch.Tensor]]:
         """
         Generate model outputs. Overridden because nnunets outputs lists when deep supervision is on so we have to
         reformat the output into dicts.
@@ -614,6 +632,7 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch
         Additionally if device type is cuda, loss computed in mixed precision.
 
         Args:
+            model (nn.Module): The model used to make predictions
             input (TorchInputType): The model inputs
 
         Returns:
@@ -624,9 +643,9 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, dict[str, torch
             # If device type is cuda, nnUNet defaults to mixed precision forward pass
             if self.device.type == "cuda":
                 with torch.autocast(self.device.type, enabled=True):
-                    output = self.model(input)
+                    output = model(input)
             else:
-                output = self.model(input)
+                output = model(input)
         else:
             raise TypeError('"input" must be of type torch.Tensor for nnUNetClient')
 
@@ -742,6 +761,15 @@ def update_metric_manager(
             target (TorchTargetType): the targets generated by the dataloader to evaluate the preds with
             metric_manager (MetricManager): the metric manager to update
         """
+
+        log(DEBUG, f"preds: {preds.keys()}")
+
+        # If personalized preds will have keys prefixed with "global-" and "local-",
+        # but we only want to keep the local ones here. NOTE: this will still work
+        # even in the non-personalized case as those preds don't get prefixed with "global-"
+        preds = {k: v for k, v in preds.items() if not k.startswith("global")}
+        preds = {k.replace("local-", ""): v for k, v in preds.items()}  # no-op in non-personalized case
+
         if len(preds) > 1:
             # for nnunet the first pred in the output list is the main one
             m_pred = convert_deep_supervision_dict_to_list(preds)[0]
@@ -914,7 +942,7 @@ def update_before_train(self, current_server_round: int) -> None:
             # freeze before the first pass, gc.collect has to check all those variables
             gc.freeze()
 
-    def transform_gradients(self, losses: TrainingLosses) -> None:
+    def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None:
         """
         Apply the gradient clipping performed by the default nnunet trainer. This is the default behavior for
         nnunet 2.5.1
@@ -922,4 +950,4 @@ def transform_gradients(self, losses: TrainingLosses) -> None:
         Args:
             losses (TrainingLosses): Not used for this transformation.
         """
-        nn.utils.clip_grad_norm_(self.model.parameters(), self.max_grad_norm)
+        nn.utils.clip_grad_norm_(model.parameters(), self.max_grad_norm)
diff --git a/fl4health/clients/perfcl_client.py b/fl4health/clients/perfcl_client.py
index 251083bd0..0b50337ee 100644
--- a/fl4health/clients/perfcl_client.py
+++ b/fl4health/clients/perfcl_client.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 
 import torch
+import torch.nn as nn
 from flwr.common.typing import Config
 
 from fl4health.checkpointing.client_module import ClientCheckpointAndStateModule
@@ -136,7 +137,7 @@ def _all_contrastive_loss_modules_defined(self) -> bool:
             and self.initial_global_module is not None
         )
 
-    def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
+    def _predict_with_model(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
         """
         Computes the prediction(s) and features of the model(s) given the input.
 
@@ -152,7 +153,7 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp
         """
         # For PerFCL models, we required the input to simply be a torch.Tensor
         assert isinstance(input, torch.Tensor)
-        preds, features = self.model(input)
+        preds, features = model(input)
         # In the first server round, these module will not have been set.
         if (
             self.old_local_module is not None
@@ -167,6 +168,22 @@ def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureTyp
 
         return preds, features
 
+    def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
+        """
+        Computes the prediction(s) and features of the model(s) given the input.
+
+        Args:
+            input (TorchInputType): Inputs to be fed into the model. ``TorchInputType`` is simply an alias
+                for the union of ``torch.Tensor`` and ``dict[str, torch.Tensor]``.
+
+        Returns:
+            tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: A tuple in which the first element
+            contains predictions indexed by name and the second element contains intermediate activations
+            index by name. Specifically the features of the model, features of the global model and features of
+            the old model are returned. All predictions included in dictionary will be used to compute metrics.
+        """
+        return self._predict_with_model(self.model, input)
+
     def update_after_train(self, local_steps: int, loss_dict: dict[str, float], config: Config) -> None:
         """
         This function is called after client-side training concludes. In this case, it is used to save the local
diff --git a/fl4health/mixins/adaptive_drift_constrained.py b/fl4health/mixins/adaptive_drift_constrained.py
index e907450f4..ea518f89c 100644
--- a/fl4health/mixins/adaptive_drift_constrained.py
+++ b/fl4health/mixins/adaptive_drift_constrained.py
@@ -16,7 +16,7 @@
 from fl4health.parameter_exchange.parameter_exchanger_base import ParameterExchanger
 from fl4health.parameter_exchange.parameter_packer import ParameterPackerAdaptiveConstraint
 from fl4health.utils.losses import TrainingLosses
-from fl4health.utils.typing import TorchFeatureType, TorchPredType, TorchTargetType
+from fl4health.utils.typing import TorchInputType, TorchPredType, TorchTargetType
 
 
 @runtime_checkable
@@ -148,39 +148,27 @@ def set_parameters(
 
         super().set_parameters(server_model_state, config, fitting_round)  # type: ignore[safe-super]
 
-    def compute_training_loss(
-        self: AdaptiveDriftConstrainedProtocol,
-        preds: TorchPredType,
-        features: TorchFeatureType,
-        target: TorchTargetType,
-    ) -> TrainingLosses:
-        """
-        Computes training loss given predictions of the model and ground truth data. Adds to objective by including
-        penalty loss.
+    def train_step(
+        self: AdaptiveDriftConstrainedProtocol, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
 
-        Args:
-            preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in
-                dictionary will be used to compute metrics.
-            features: (TorchFeatureType): Feature(s) of the model(s) indexed by name.
-            target: (TorchTargetType): Ground truth data to evaluate predictions against.
+        losses, preds = self._compute_preds_and_losses(self.model, self.optimizers["global"], input, target)
+        loss_clone = losses.backward["backward"].clone()
 
-        Returns:
-            TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed
-            by name. Additional losses includes penalty loss.
-        """
-        loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target)
-        if additional_losses is None:
-            additional_losses = {}
-
-        additional_losses["loss"] = loss.clone()
-        # adding the vanilla loss to the additional losses to be used by update_after_train for potential adaptation
-        additional_losses["loss_for_adaptation"] = loss.clone()
-
-        # Compute the drift penalty loss and store it in the additional losses dictionary.
+        # apply penalty
         penalty_loss = self.compute_penalty_loss()
-        additional_losses["penalty_loss"] = penalty_loss.clone()
-
-        return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses)
+        losses.backward["backward"] = losses.backward["backward"] + penalty_loss
+        losses = self._apply_backwards_on_losses_and_take_step(self.model, self.optimizers["global"], losses)
+
+        # prepare return values
+        additional_losses = {
+            "penalty_loss": penalty_loss.clone(),
+            "local_loss": loss_clone,
+            "loss_for_adaptation": loss_clone.clone(),
+        }
+        losses.additional_losses = additional_losses
+
+        return losses, preds
 
     def get_parameter_exchanger(self: AdaptiveDriftConstrainedProtocol, config: Config) -> ParameterExchanger:
         """
diff --git a/fl4health/mixins/core_protocols.py b/fl4health/mixins/core_protocols.py
index d28ff7794..6d044edcb 100644
--- a/fl4health/mixins/core_protocols.py
+++ b/fl4health/mixins/core_protocols.py
@@ -77,12 +77,41 @@ def initialize_all_model_weights(self, parameters: NDArrays, config: Config) ->
     def update_before_train(self, current_server_round: int) -> None:
         pass  # pragma: no cover
 
+    def _compute_preds_and_losses(
+        self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        pass  # pragma: no cover
+
+    def _apply_backwards_on_losses_and_take_step(
+        self, model: nn.Module, optimizer: Optimizer, losses: TrainingLosses
+    ) -> TrainingLosses:
+        pass  # pragma: no cover
+
+    def _train_step_with_model_and_optimizer(
+        self, model: nn.Module, optimizer: Optimizer, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[TrainingLosses, TorchPredType]:
+        pass  # pragma: no cover
+
+    def _val_step_with_model(
+        self, model: nn.Module, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[EvaluationLosses, TorchPredType]:
+        pass  # pragma: no cover
+
+    def _predict_with_model(self, model: nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
+        pass  # pragma: no cover
+
     def predict(self, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
         pass  # pragma: no cover
 
     def transform_target(self, target: TorchTargetType) -> TorchTargetType:
         pass  # pragma: no cover
 
+    def _transform_gradients_with_model(self, model: torch.nn.Module, losses: TrainingLosses) -> None:
+        pass  # pragma: no cover
+
+    def transform_gradients(self, losses: TrainingLosses) -> None:
+        pass  # pragma: no cover
+
     def compute_training_loss(
         self,
         preds: TorchPredType,
diff --git a/fl4health/mixins/personalized/ditto.py b/fl4health/mixins/personalized/ditto.py
index 7ceae19c7..6137aa81c 100644
--- a/fl4health/mixins/personalized/ditto.py
+++ b/fl4health/mixins/personalized/ditto.py
@@ -1,7 +1,8 @@
 """Ditto Personalized Mixin"""
 
+import copy
 import warnings
-from logging import DEBUG, INFO, WARN
+from logging import INFO, WARN
 from typing import Any, Protocol, cast, runtime_checkable
 
 import torch
@@ -34,9 +35,6 @@ def _copy_optimizer_with_new_params(self, original_optimizer: Optimizer) -> Opti
     def set_initial_global_tensors(self) -> None:
         pass  # pragma: no cover
 
-    def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        pass  # pragma: no cover
-
     def safe_global_model(self) -> nn.Module:
         pass  # pragma: no cover
 
@@ -97,10 +95,7 @@ def __init_subclass__(cls, **kwargs: Any) -> None:
             f"base classes implement BasicClient. This may cause runtime errors."
         )
         log(WARN, msg)
-        warnings.warn(
-            msg,
-            RuntimeWarning,
-        )
+        warnings.warn(msg, RuntimeWarning)
 
     def safe_global_model(self: DittoPersonalizedProtocol) -> nn.Module:
         """Convenient accessor for the global model.
@@ -129,7 +124,7 @@ def _copy_optimizer_with_new_params(self: DittoPersonalizedProtocol, original_op
         """Helper method to make a copy of the original optimizer for the global model.
 
         Args:
-            original_optimizer (Optimizer): original optimizer of the underyling `BasicClient`.
+            original_optimizer (Optimizer): original optimizer of the underlying `BasicClient`.
 
         Returns:
             Optimizer: a copy of the original optimizer to be used by the global model.
@@ -175,7 +170,7 @@ def get_global_model(self: DittoPersonalizedProtocol, config: Config) -> nn.Modu
         Returns:
             nn.Module: The PyTorch model serving as the global model for Ditto
         """
-        return self.get_model(config).to(self.device)
+        return copy.deepcopy(self.get_model(config).to(self.device))
 
     @ensure_protocol_compliance
     def get_optimizer(self: DittoPersonalizedProtocol, config: Config) -> dict[str, Optimizer]:
@@ -369,182 +364,62 @@ def train_step(
             corresponding to predictions from the global and local Ditto models for metric evaluations.
         """
 
-        # Clear gradients from optimizers if they exist
-        self.optimizers["global"].zero_grad()
-        self.optimizers["local"].zero_grad()
-
-        # Forward pass on both the global and local models
-        preds, features = self.predict(input)
-        target = self.transform_target(target)  # Apply transformation (Defaults to identity)
-
-        # Compute all relevant losses
-        losses = self.compute_training_loss(preds, features, target)
-
-        # Take a step with the global model vanilla loss
-        losses.additional_losses["global_loss"].backward()
-        self.optimizers["global"].step()
-
-        # Take a step with the local model using the local loss and Ditto constraint
-        losses.backward["backward"].backward()
-        self.optimizers["local"].step()
-
-        # Return dictionary of predictions where key is used to name respective MetricMeters
-        return losses, preds
-
-    def predict(
-        self: DittoPersonalizedProtocol,
-        input: TorchInputType,
-    ) -> tuple[TorchPredType, TorchFeatureType]:
-        """
-        Computes the predictions for both the **GLOBAL** and **LOCAL** models and pack them into the prediction
-        dictionary
-
-        Args:
-            input (TorchInputType): Inputs to be fed into both models.
-
-        Returns:
-            tuple[TorchPredType, TorchFeatureType]: A tuple in which the first element contains predictions indexed by
-            name and the second element contains intermediate activations index by name. For Ditto, we only need the
-            predictions, so the second dictionary is simply empty.
-
-        Raises:
-            ValueError: Occurs when something other than a tensor or dict of tensors is returned by the model
-                forward.
-        """
+        # global
+        global_losses, global_preds = self._compute_preds_and_losses(
+            self.safe_global_model(), self.optimizers["global"], input, target
+        )
+        # local
+        local_losses, local_preds = self._compute_preds_and_losses(self.model, self.optimizers["local"], input, target)
+        local_loss_clone = local_losses.backward["backward"].clone()  # need a clone for later
+
+        # take step
+        # global
+        global_losses = self._apply_backwards_on_losses_and_take_step(
+            self.safe_global_model(), self.optimizers["global"], global_losses
+        )
+        # local
+        penalty_loss = self.compute_penalty_loss()
+        local_losses.backward["backward"] = local_losses.backward["backward"] + penalty_loss
+        local_losses = self._apply_backwards_on_losses_and_take_step(
+            self.model, self.optimizers["local"], local_losses
+        )
 
-        if hasattr(self, "_predict"):
-            log(INFO, "Using '_predict' to make predictions")
-            global_preds, _ = self._predict(self.safe_global_model(), input)
-            local_preds, _ = self._predict(self.model, input)
-            log(INFO, "Successfully predicted for global and local models")
-        else:
-            if isinstance(input, torch.Tensor):
-                global_preds = self.safe_global_model()(input)
-                local_preds = self.model(input)
-            elif isinstance(input, dict):
-                # If input is a dictionary, then we unpack it before computing the forward pass.
-                # Note that this assumes the keys of the input match (exactly) the keyword args
-                # of the forward method.
-                global_preds = self.safe_global_model()(**input)
-                local_preds = self.model(**input)
+        # prepare return values
+        additional_losses = {
+            "penalty_loss": penalty_loss.clone(),
+            "local_loss": local_loss_clone,
+            "global_loss": global_losses.backward["backward"],
+            "loss_for_adaptation": local_loss_clone.clone(),
+        }
+        local_losses.additional_losses = additional_losses
 
+        # combined preds
         if isinstance(global_preds, torch.Tensor) and isinstance(local_preds, torch.Tensor):
-            return {"global": global_preds, "local": local_preds}, {}
+            combined_preds = {"global": global_preds, "local": local_preds}
         elif isinstance(global_preds, dict) and isinstance(local_preds, dict):
-            retval = {f"global-{k}": v for k, v in global_preds.items()}
-            retval.update(**{f"local-{k}": v for k, v in local_preds.items()})
-            return retval, {}
-        else:
-            raise ValueError(f"Unsupported pred types. Global: {type(global_preds)}, local: {type(local_preds)}")
-
-    def _extract_pred(self, kind: str, preds: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
-        """Helper method to extract predictions from global and local models.
-
-        Args:
-            kind (str): the kind of predictions to be extracted i.e., for 'local' or for 'global'
-            preds (dict[str, torch.Tensor]): the full preds result from `self.predict()`
-
-        Raises:
-            ValueError: If supplied `kind` is not of 'global' or 'local'
-
-        Returns:
-            dict[str, torch.Tensor]: the predictions associated with the specified kind
-        """
-        if kind not in ["global", "local"]:
-            raise ValueError("Unsupported kind of prediction. Must be 'global' or 'local'.")
-
-        # filter
-        retval = {k: v for k, v in preds.items() if kind in k}
-        # remove prefix
-        retval = {k.replace(f"{kind}-", ""): v for k, v in retval.items()}
-        return retval
-
-    def compute_loss_and_additional_losses(
-        self: DittoPersonalizedProtocol,
-        preds: TorchPredType,
-        features: TorchFeatureType,
-        target: TorchTargetType,
-    ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
-        """
-        Computes the local model loss and the global Ditto model loss (stored in additional losses) for reporting and
-        training of the global model
-
-        Args:
-            preds (TorchPredType): Prediction(s) of the model(s) indexed by name.
-            features (TorchFeatureType): Feature(s) of the model(s) indexed by name.
-            target (TorchTargetType): Ground truth data to evaluate predictions against.
-        Returns:
-            tuple[torch.Tensor, dict[str, torch.Tensor]]: A tuple with:
-
-            - The tensor for the model loss
-            - A dictionary with ``local_loss``, ``global_loss`` as additionally reported loss values.
-        """
-
-        global_preds = self._extract_pred(kind="global", preds=preds)
-        local_preds = self._extract_pred(kind="local", preds=preds)
-
-        log(DEBUG, f"global_preds has keys {global_preds.keys()}")
-        log(DEBUG, f"local_preds has keys {local_preds.keys()}")
-
-        # Compute global model vanilla loss
-
-        if hasattr(self, "_special_compute_loss_and_additional_losses"):
-            log(INFO, "Using '_special_compute_loss_and_additional_losses' to compute loss")
-            global_loss, _ = self._special_compute_loss_and_additional_losses(global_preds, features, target)
-
-            # Compute local model loss
-            local_loss, _ = self._special_compute_loss_and_additional_losses(local_preds, features, target)
-
-        else:
-            global_loss = self.criterion(global_preds["global"], target)
-
-            # Compute local model loss
-            local_loss = self.criterion(local_preds["local"], target)
-
-        additional_losses = {"local_loss": local_loss.clone(), "global_loss": global_loss}
+            combined_preds = {f"global-{k}": v for k, v in global_preds.items()}
+            combined_preds.update(**{f"local-{k}": v for k, v in local_preds.items()})
 
-        return local_loss, additional_losses
+        return local_losses, combined_preds
 
-    def compute_training_loss(
-        self: DittoPersonalizedProtocol,
-        preds: TorchPredType,
-        features: TorchFeatureType,
-        target: TorchTargetType,
-    ) -> TrainingLosses:
-        """
-        Computes training losses given predictions of the global and local models and ground truth data.
-        For the local model we add to the vanilla loss function by including Ditto penalty loss which is the l2 inner
-        product between the initial global model weights and weights of the local model. This is stored in backward
-        The loss to optimize the global model is stored in the additional losses dictionary under "global_loss"
-
-        Args:
-            preds (TorchPredType): Prediction(s) of the model(s) indexed by name. All predictions included in
-                dictionary will be used to compute metrics.
-            features: (TorchFeatureType): Feature(s) of the model(s) indexed by name.
-            target: (TorchTargetType): Ground truth data to evaluate predictions against.
-
-        Returns:
-            TrainingLosses: An instance of ``TrainingLosses`` containing backward loss and additional losses indexed by
-            name. Additional losses includes each loss component and the global model
-            loss tensor.
-        """
-        # Check that both models are in training mode
-        assert self.safe_global_model().training and self.model.training
-
-        # local loss is stored in loss, global model loss is stored in additional losses.
-        loss, additional_losses = self.compute_loss_and_additional_losses(preds, features, target)
-        additional_losses = additional_losses or {}  # make mypy happy
-
-        # Setting the adaptation loss to that of the local model, as its performance should dictate whether more or
-        # less weight is used to constrain it to the global model (as in FedProx)
-        additional_losses["loss_for_adaptation"] = additional_losses["local_loss"].clone()
+    def val_step(
+        self: DittoPersonalizedProtocol, input: TorchInputType, target: TorchTargetType
+    ) -> tuple[EvaluationLosses, TorchPredType]:
 
-        # This is the Ditto penalty loss of the local model compared with the original Global model weights, scaled
-        # by drift_penalty_weight (or lambda in the original paper)
-        penalty_loss = self.compute_penalty_loss()
-        additional_losses["penalty_loss"] = penalty_loss.clone()
+        # global
+        global_losses, global_preds = self._val_step_with_model(self.safe_global_model(), input, target)
+        # local
+        local_losses, local_preds = self._val_step_with_model(self.model, input, target)
 
-        return TrainingLosses(backward=loss + penalty_loss, additional_losses=additional_losses)
+        # combine
+        losses = EvaluationLosses(
+            local_losses.checkpoint,
+            additional_losses={"global_loss": global_losses.checkpoint, "local_loss": local_losses.checkpoint},
+        )
+        preds: TorchPredType = {}
+        preds.update(**{f"global-{k}": v for k, v in global_preds.items()})
+        preds.update(**{f"local-{k}": v for k, v in local_preds.items()})
+        return losses, preds
 
     @ensure_protocol_compliance
     def validate(
diff --git a/fl4health/strategies/fedavg_with_adaptive_constraint.py b/fl4health/strategies/fedavg_with_adaptive_constraint.py
index 28f0f1370..2fb1227fb 100644
--- a/fl4health/strategies/fedavg_with_adaptive_constraint.py
+++ b/fl4health/strategies/fedavg_with_adaptive_constraint.py
@@ -28,7 +28,7 @@ def __init__(
         on_fit_config_fn: Callable[[int], dict[str, Scalar]] | None = None,
         on_evaluate_config_fn: Callable[[int], dict[str, Scalar]] | None = None,
         accept_failures: bool = True,
-        initial_parameters: Parameters,
+        initial_parameters: Parameters | None,
         fit_metrics_aggregation_fn: MetricsAggregationFn | None = None,
         evaluate_metrics_aggregation_fn: MetricsAggregationFn | None = None,
         initial_loss_weight: float = 1.0,
@@ -102,8 +102,9 @@ def __init__(
 
         self.previous_loss = float("inf")
 
-        self.server_model_weights = parameters_to_ndarrays(initial_parameters)
-        initial_parameters.tensors.extend(ndarrays_to_parameters([np.array(initial_loss_weight)]).tensors)
+        if initial_parameters:
+            self.server_model_weights = parameters_to_ndarrays(initial_parameters)
+            initial_parameters.tensors.extend(ndarrays_to_parameters([np.array(initial_loss_weight)]).tensors)
 
         super().__init__(
             fraction_fit=fraction_fit,
diff --git a/tests/mixins/personalized/test_ditto.py b/tests/mixins/personalized/test_ditto.py
index e4b464ca9..52f4a59a3 100644
--- a/tests/mixins/personalized/test_ditto.py
+++ b/tests/mixins/personalized/test_ditto.py
@@ -11,7 +11,8 @@
 from numpy.testing import assert_array_equal
 from torch.nn.modules.loss import _Loss
 from torch.optim import Optimizer
-from torch.testing import assert_close
+
+# from torch.testing import assert_close
 from torch.utils.data import DataLoader, TensorDataset
 
 from fl4health.clients.basic_client import BasicClient
@@ -27,7 +28,7 @@
 from fl4health.parameter_exchange.parameter_packer import (
     ParameterPackerAdaptiveConstraint,
 )
-from fl4health.utils.losses import TrainingLosses
+from fl4health.utils.losses import EvaluationLosses, TrainingLosses
 from fl4health.utils.typing import TorchFeatureType, TorchInputType, TorchPredType
 
 
@@ -58,7 +59,9 @@ def get_optimizer(self, config: Config) -> Optimizer | dict[str, Optimizer]:
     def get_criterion(self, config: Config) -> _Loss:
         return torch.nn.CrossEntropyLoss()
 
-    def _predict(self, model: torch.nn.Module, input: TorchInputType) -> tuple[TorchPredType, TorchFeatureType]:
+    def _predict_with_model(
+        self, model: torch.nn.Module, input: TorchInputType
+    ) -> tuple[TorchPredType, TorchFeatureType]:
         return {}, {}
 
 
@@ -153,6 +156,74 @@ def test_get_parameters() -> None:
     assert_array_equal(packed_params, pack_params_return_val)
 
 
+@patch.object(_TestDittoedClient, "compute_penalty_loss")
+@patch.object(_TestDittoedClient, "_apply_backwards_on_losses_and_take_step")
+@patch.object(_TestDittoedClient, "_compute_preds_and_losses")
+def test_predict(
+    mock_private_compute_preds_and_losses: MagicMock,
+    mock_private_apply_backwards_on_losses_and_take_step: MagicMock,
+    mock_compute_penalty_loss: MagicMock,
+) -> None:
+    # setup client
+    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+    client.model = torch.nn.Linear(5, 5)
+    client.global_model = torch.nn.Linear(5, 5)
+    client.optimizers = {
+        "global": torch.optim.SGD(client.model.parameters(), lr=0.0001),
+        "local": torch.optim.SGD(client.model.parameters(), lr=0.0001),
+    }
+    # setup mocks
+    mock_param_exchanger = MagicMock()
+    push_params_return_model_weights: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    pack_params_return_val: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    mock_param_exchanger.push_parameters.return_value = push_params_return_model_weights
+    mock_param_exchanger.pack_parameters.return_value = pack_params_return_val
+    client.parameter_exchanger = mock_param_exchanger
+    client.initialized = True
+
+    mock_backward_loss = MagicMock()
+    mock_additional_global_loss = MagicMock()
+    dummy_training_losses = TrainingLosses(
+        backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss}
+    )
+    dummy_training_losses_for_local = TrainingLosses(
+        backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss}
+    )
+    mock_private_compute_preds_and_losses.side_effect = [
+        (
+            dummy_training_losses,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+        (
+            dummy_training_losses_for_local,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+    ]
+    mock_private_apply_backwards_on_losses_and_take_step.side_effect = [
+        dummy_training_losses,
+        dummy_training_losses_for_local,
+    ]
+
+    # act
+    input, target = torch.tensor([1, 1, 1, 1, 1]), torch.zeros(3)
+    # TODO: fix the mixin/protocol typing that leads to mypy complaint
+    _ = client.train_step(input, target)  # type: ignore
+
+    mock_private_compute_preds_and_losses.assert_has_calls(
+        [
+            _Call(((client.global_model, client.optimizers["global"], input, target), {})),
+            _Call(((client.model, client.optimizers["local"], input, target), {})),
+        ]
+    )
+    mock_private_apply_backwards_on_losses_and_take_step.assert_has_calls(
+        [
+            _Call(((client.global_model, client.optimizers["global"], dummy_training_losses), {})),
+            _Call(((client.model, client.optimizers["local"], dummy_training_losses_for_local), {})),
+        ]
+    )
+    mock_compute_penalty_loss.assert_called_once()
+
+
 @patch.object(_TestDittoedClient, "setup_client")
 @patch("fl4health.mixins.personalized.ditto.FullParameterExchanger")
 def test_get_parameters_uninitialized(mock_param_exchanger: MagicMock, mock_setup_client: MagicMock) -> None:
@@ -239,92 +310,6 @@ def test_get_optimizer(mock_copy_optimizer: MagicMock) -> None:
     mock_copy_optimizer.assert_called_once_with(client.optimizers["local"])
 
 
-def test_predict() -> None:
-    # setup client
-    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-
-    mock_model = MagicMock()
-    mock_global_model = MagicMock()
-
-    mock_model.return_value = torch.ones(5)
-    mock_global_model.return_value = torch.zeros(5)
-
-    client.model = mock_model
-    client.global_model = mock_global_model
-
-    client.optimizers = {
-        "global": MagicMock(),
-        "local": MagicMock(),
-    }
-
-    client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint())
-    client.initialized = True
-
-    # act
-    # TODO: fix the mixin/protocol typing that leads to mypy complaint
-    res, _ = client.predict(input=torch.zeros(5))  # type: ignore
-    print(f"res: {res}")
-    print(torch.zeros(5))
-
-    # assert
-    assert_close(res["global"], torch.zeros(5))
-    assert_close(res["local"], torch.ones(5))
-
-
-@patch.object(_TestDittoedClientV2, "_predict")
-def test_predict_delagation(private_predict: MagicMock) -> None:
-    # setup client
-    client = _TestDittoedClientV2(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-    client.model = torch.nn.Linear(5, 5)
-    client.global_model = torch.nn.Linear(5, 5)
-
-    private_predict.side_effect = [(torch.zeros(5), {}), (torch.ones(5), {})]
-
-    client.optimizers = {
-        "global": MagicMock(),
-        "local": MagicMock(),
-    }
-
-    client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint())
-    client.initialized = True
-
-    # act
-    # TODO: fix the mixin/protocol typing that leads to mypy complaint
-    res, _ = client.predict(input=torch.zeros(5))  # type: ignore
-    print(f"res: {res}")
-    print(torch.zeros(5))
-
-    # assert
-    assert_close(res["global"], torch.zeros(5))
-    assert_close(res["local"], torch.ones(5))
-
-
-def test_extract_pred() -> None:
-    # setup client
-    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-
-    res = client._extract_pred(
-        kind="global", preds={"global-xyz": torch.ones(5), "global-abc": torch.zeros(5), "local": torch.zeros(5)}
-    )
-
-    assert_close(res["xyz"], torch.ones(5))
-    assert_close(res["abc"], torch.zeros(5))
-
-
-def test_extract_pred_raises_error() -> None:
-    # setup client
-    client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
-
-    with pytest.raises(ValueError):
-        client._extract_pred(
-            kind="oops", preds={"global-xyz": torch.ones(5), "global-abc": torch.zeros(5), "local": torch.zeros(5)}
-        )
-
-
 @patch.object(_TestDittoedClient, "set_initial_global_tensors")
 @patch.object(_TestBasicClient, "update_before_train")
 def test_update_before_train(
@@ -350,43 +335,52 @@ def test_safe_model_raises_error() -> None:
         client.safe_global_model()  # type: ignore
 
 
-@patch("torch.optim.Optimizer")
-@patch.object(_TestDittoedClient, "predict")
-@patch.object(_TestDittoedClient, "compute_training_loss")
-def test_train_step(
-    mock_compute_training_loss: MagicMock, mock_predict: MagicMock, mock_optimizer_class: MagicMock
+@patch.object(_TestDittoedClient, "_val_step_with_model")
+def test_val_step(
+    mock_private_val_step_with_model: MagicMock,
 ) -> None:
     # setup client
     client = _TestDittoedClient(data_path=Path(""), metrics=[Accuracy()], device=torch.device("cpu"))
+    client.model = torch.nn.Linear(5, 5)
+    client.global_model = torch.nn.Linear(5, 5)
     client.optimizers = {
-        "global": MagicMock(),
-        "local": MagicMock(),
+        "global": torch.optim.SGD(client.model.parameters(), lr=0.0001),
+        "local": torch.optim.SGD(client.model.parameters(), lr=0.0001),
     }
-    client.train_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.val_loader = DataLoader(TensorDataset(torch.ones((1000, 28, 28, 1)), torch.ones((1000))))
-    client.parameter_exchanger = FullParameterExchangerWithPacking(ParameterPackerAdaptiveConstraint())
+    # setup mocks
+    mock_param_exchanger = MagicMock()
+    push_params_return_model_weights: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    pack_params_return_val: NDArray = np.ndarray(shape=(2, 2), dtype=float)
+    mock_param_exchanger.push_parameters.return_value = push_params_return_model_weights
+    mock_param_exchanger.pack_parameters.return_value = pack_params_return_val
+    client.parameter_exchanger = mock_param_exchanger
     client.initialized = True
 
-    # arrange mocks
-    pred_return: tuple[TorchPredType, TorchFeatureType] = {"local": torch.ones(5)}, {}
-    preds, features = pred_return
-    mock_predict.return_value = preds, features
-    mock_backward_loss = MagicMock()
-    mock_additional_global_loss = MagicMock()
-    losses = TrainingLosses(
-        backward={"backward": mock_backward_loss}, additional_losses={"global_loss": mock_additional_global_loss}
+    dummy_training_losses = EvaluationLosses(
+        checkpoint=torch.ones(5),
     )
-    mock_compute_training_loss.return_value = losses
+    dummy_training_losses_for_local = EvaluationLosses(
+        checkpoint=torch.ones(5),
+    )
+    mock_private_val_step_with_model.side_effect = [
+        (
+            dummy_training_losses,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+        (
+            dummy_training_losses_for_local,
+            {"prediction": torch.Tensor([1, 2, 3, 4, 5])},
+        ),
+    ]
 
     # act
     input, target = torch.tensor([1, 1, 1, 1, 1]), torch.zeros(3)
     # TODO: fix the mixin/protocol typing that leads to mypy complaint
-    retval = client.train_step(input, target)  # type: ignore
-
-    mock_predict.assert_called_once()
-    client.optimizers["global"].zero_grad.assert_called_once()
-    client.optimizers["local"].zero_grad.assert_called_once()
-    mock_compute_training_loss.assert_called_once_with(preds, features, target)
-    mock_backward_loss.backward.assert_called_once()
-    mock_additional_global_loss.backward.assert_called_once()
-    assert retval == (losses, preds)
+    _ = client.val_step(input, target)  # type: ignore
+
+    mock_private_val_step_with_model.assert_has_calls(
+        [
+            _Call(((client.global_model, input, target), {})),
+            _Call(((client.model, input, target), {})),
+        ]
+    )