The error while using mean function in tensorlow

Question

I am facing following problem while running the code. Its might be a minor, but I am new to python. I would appreciate correction.

import keras
import numpy as np
import tensorflow as tf
from tensorflow.python.ops import math_ops    


def custom_loss(y_true, y_pred):
  diff = math_ops.squared_difference(y_pred, y_true)  #squared difference
  print(diff)
  loss = Mean(diff, axis=-1) #mean over last dimension
  loss = loss / 10.0
return loss

 model = keras.Sequential([
               keras.layers.Dense(10, activation='relu', input_shape=(1,)),
               keras.layers.Dense(1)
 ])

 model.compile(loss=custom_loss, optimizer='sgd')

 X_train = np.array([[10.0],[20.0], [30.0],[40.0],[50.0],[60.0],[10.0], [20.0]])
 y_train = np.array([6.0, 12, 18,24,30, 36,6, 12]) #dummy data

 model.fit(X_train, y_train, batch_size=2, epochs=10)

Following error is occurring. Please let me know how to correct it. Regards

NameError: in user code:

File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1021, in train_function  *
    return step_function(self, iterator)
File "<ipython-input-6-a658c19a56c4>", line 10, in custom_loss  *
    loss = Mean(diff, axis=-1) #mean over last dimension

NameError: name 'Mean' is not defined