Custom MSE loss for tensorflow

NMSE.py

# Importing the package
import tensorflow as tf


class NMSE(tf.keras.losses.Loss):
    def __init__(self):
        super().__init__()

    def call(self, y_true, y_pred):
        # Calculating the mse; 
        # Adding an additional division by 2 to mimic sklearn
        return tf.reduce_mean(tf.square(y_true - y_pred)) / 2