def METHOD_NAME(self, keras_model, compression_loss_fn):
    filter_fn = self.make_filter_trainable_variables_fn()
    trainable_variables = filter_fn(keras_model.trainable_variables)
    def _total_loss_fn(labels, outputs):
        cls_loss = self._cls_loss_fn(outputs["cls_outputs"], labels["cls_targets"], labels["num_positives"])
        box_loss = self._box_loss_fn(outputs["box_outputs"], labels["box_targets"], labels["num_positives"])
        model_loss = cls_loss + self._box_loss_weight * box_loss
        l2_regularization_loss = self.weight_decay_loss(trainable_variables)
        compression_loss = compression_loss_fn()
        total_loss = model_loss + l2_regularization_loss + compression_loss
        return {
            "total_loss": total_loss,
            "cls_loss": cls_loss,
            "box_loss": box_loss,
            "model_loss": model_loss,
            "l2_regularization_loss": l2_regularization_loss,
            "compression_loss": compression_loss,
        }
    return _total_loss_fn