def METHOD_NAME(
    self,
    shareable: Shareable,
    fl_ctx: FLContext,
    abort_signal: Signal,
) -> Shareable:
    """Training task pipeline for Ditto
    Get global model weights (potentially with HE)
    Prepare for fedprox loss
    Load Ditto personalized model info
    Local training reference model and personalized model
    Return updated weights of reference model (model_diff)
    """
    if abort_signal.triggered:
        return make_reply(ReturnCode.TASK_ABORTED)
    # get round information
    current_round = shareable.get_header(AppConstants.CURRENT_ROUND)
    total_rounds = shareable.get_header(AppConstants.NUM_ROUNDS)
    self.log_info(fl_ctx, f"Current/Total Round: {current_round + 1}/{total_rounds}")
    self.log_info(fl_ctx, f"Client identity: {fl_ctx.get_identity_name()}")
    # update local model weights with received weights
    dxo = from_shareable(shareable)
    global_weights = dxo.data
    # Before loading weights, tensors might need to be reshaped to support HE for secure aggregation.
    local_var_dict = self.model.state_dict()
    model_keys = global_weights.keys()
    for var_name in local_var_dict:
        if var_name in model_keys:
            weights = global_weights[var_name]
            try:
                # reshape global weights to compute difference later on
                global_weights[var_name] = np.reshape(weights, local_var_dict[var_name].shape)
                # update the local dict
                local_var_dict[var_name] = torch.as_tensor(global_weights[var_name])
            except Exception as e:
                raise ValueError("Convert weight from {} failed with error: {}".format(var_name, str(e)))
    self.model.load_state_dict(local_var_dict)
    # Load Ditto personalized model
    self.ditto_helper.load_model(local_var_dict)
    # local steps
    epoch_len = len(self.train_loader)
    self.log_info(fl_ctx, f"Local steps per epoch: {epoch_len}")
    # make a copy of model_global as reference for
    # 1. FedProx loss of reference model
    # 2. Ditto loss of personalized model
    model_global = copy.deepcopy(self.model)
    for param in model_global.parameters():
        param.requires_grad = False
    # local train reference model
    self.local_train(
        fl_ctx=fl_ctx,
        train_loader=self.train_loader,
        model_global=model_global,
        abort_signal=abort_signal,
    )
    if abort_signal.triggered:
        return make_reply(ReturnCode.TASK_ABORTED)
    self.epoch_of_start_time += self.aggregation_epochs
    # local train ditto model
    self.ditto_helper.local_train(
        train_loader=self.train_loader,
        model_global=model_global,
        abort_signal=abort_signal,
        writer=self.writer,
    )
    if abort_signal.triggered:
        return make_reply(ReturnCode.TASK_ABORTED)
    # local valid ditto model each round
    metric = self.local_valid(
        self.ditto_helper.model,
        self.valid_loader,
        abort_signal,
        tb_id="val_metric_per_model",
        record_epoch=self.ditto_helper.epoch_global,
    )
    if abort_signal.triggered:
        return make_reply(ReturnCode.TASK_ABORTED)
    self.log_info(fl_ctx, f"val_metric_per_model: {metric:.4f}")
    # save model
    self.ditto_helper.update_metric_save_model(metric=metric)
    # compute delta model, global model has the primary key set
    local_weights = self.model.state_dict()
    model_diff = {}
    for name in global_weights:
        if name not in local_weights:
            continue
        model_diff[name] = np.subtract(local_weights[name].cpu().numpy(), global_weights[name], dtype=np.float32)
        if np.any(np.isnan(model_diff[name])):
            self.system_panic(f"{name} weights became NaN...", fl_ctx)
            return make_reply(ReturnCode.EXECUTION_EXCEPTION)
    # flush the tb writer
    self.writer.flush()
    # build the shareable
    dxo = DXO(data_kind=DataKind.WEIGHT_DIFF, data=model_diff)
    dxo.set_meta_prop(MetaKey.NUM_STEPS_CURRENT_ROUND, epoch_len)
    self.log_info(fl_ctx, "Local epochs finished. Returning shareable")
    return dxo.to_shareable()