如何将gpflow GPR的批量训练编译成tf.function？

import tensorflow as tf
import gpflow
from gpflow.models import GPR

def batch_train(X, Y):
    # 数据准备
    X_tensor = tf.convert_to_tensor(X)
    Y_tensor = tf.convert_to_tensor(Y)
    
    # 创建 GPR 模型
    kernel = gpflow.kernels.RBF()
    model = GPR(data=(X_tensor, Y_tensor), kernel=kernel)
    
    # 使用最大似然估计进行训练
    opt = gpflow.optimizers.Scipy()
    opt.minimize(model.training_loss, model.trainable_variables)
    
    return model

batch_train_tf = tf.function(batch_train)

model = batch_train_tf(X_data, Y_data)