原代码
with tf.GradientTape() as tape:
# 打平操作，[b, 28, 28] => [b, 784]
x = tf.reshape(x, (-1, 28*28))
# Step1. 得到模型输出output [b, 784] => [b, 10]
out = model(x)
# [b] => [b, 10]
y_onehot = tf.one_hot(y, depth=10)
# 计算差的平方和，[b, 10]
loss = tf.square(out-y_onehot)
# 计算每个样本的平均误差，[b] x.shape[0]相当于N
loss = tf.reduce_sum(loss) / x.shape[0]
错图:
错因：除数和被除数的数据类型不一致
解决方式:将x.shape[0]通过tf.cast改变数据类型即可
- loss = tf.reduce_sum(loss) / x.shape
改为
X = tf.cast(x.shape[0],dtype=float)
loss = tf.reduce_sum(loss) / X
觉得有用请留下你的赞 ^__^
