def METHOD_NAME(inputs, vec):
    y, z_inputs = self.inverse_mvp.predict_step((inputs, vec)).output
    jac_y, _ = self.mlp.compute_vjp(z_inputs, y)
    jac_y = torch.cat(
        (jac_y, torch.zeros(jac_y.shape[0], output_dim - input_dim)),
        dim=-1)
    jac_y += fullrank_diag_weight * y
    loss = torch.nn.functional.mse_loss(jac_y, vec)
    return loss