def METHOD_NAME(self, keras_model, compression_loss_fn):
    def _total_loss_fn(labels, outputs):
        loss_fn_out = self._loss_fn(
            labels,
            outputs,
            self._params.anchors,
            self._params.model_params.num_classes,
            ignore_thresh=self._params.model_params.loss_params.ignore_thresh,
            label_smoothing=self._params.model_params.loss_params.label_smoothing,
            elim_grid_sense=self._params.elim_grid_sense,
            use_focal_loss=self._params.model_params.loss_params.use_focal_loss,
            use_focal_obj_loss=self._params.model_params.loss_params.use_focal_obj_loss,
            use_softmax_loss=self._params.model_params.loss_params.use_softmax_loss,
            use_giou_loss=self._params.model_params.loss_params.use_giou_loss,
            use_diou_loss=self._params.model_params.loss_params.use_diou_loss,
        )
        loss, total_location_loss, total_confidence_loss, total_class_loss = loss_fn_out
        compression_loss = compression_loss_fn()
        loss += compression_loss
        return {
            "total_loss": loss,
            "total_location_loss": total_location_loss,
            "total_confidence_loss": total_confidence_loss,
            "total_class_loss": total_class_loss,
            "compression_loss": compression_loss,
        }
    return _total_loss_fn