Update out-of-date Loss example in docstring.
@@ -44,9 +44,7 @@ class Loss:
class MeanSquaredError(Loss):
def call(self, y_true, y_pred):
- y_pred = tf.convert_to_tensor_v2(y_pred)
- y_true = tf.cast(y_true, y_pred.dtype)
- return tf.reduce_mean(math_ops.square(y_pred - y_true), axis=-1)
+ return tf.reduce_mean(tf.math.square(y_pred - y_true), axis=-1)
```
When used with `tf.distribute.Strategy`, outside of built-in training loops
