如何使用TensorFlow2.0解决线性回归代码中的这个错误？

import tensorflow as tf
import numpy as np

# 定义输入特征和标签
x_train = np.array([1, 2, 3, 4, 5], dtype=np.float32)
y_train = np.array([2, 4, 6, 8, 10], dtype=np.float32)

# 定义单层线性回归模型
class LinearRegression(tf.keras.Model):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.dense = tf.keras.layers.Dense(units=1)
    
    def call(self, inputs):
        return self.dense(inputs)
    
model = LinearRegression()

# 定义损失函数和优化器
loss_object = tf.keras.losses.MeanSquaredError()
optimizer = tf.keras.optimizers.SGD(learning_rate=0.01)

# 定义训练步骤
@tf.function
def train_step(inputs, labels):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = loss_object(labels, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

# 进行多轮训练
epochs = 100
for epoch in range(epochs):
    train_step(x_train, y_train)