def METHOD_NAME(self, y_gt, y_pr, sample_weight=None):
    shape = tf.shape(y_gt)
    n_samples = shape[0]
    channel_mse = tf.reduce_sum(
        tf.math.squared_difference(
            tf.gather(y_gt, self.channel_ind, axis=3),
            tf.gather(y_pr, self.channel_ind, axis=3),
        )
    )  # rank 4
    self.height.assign(shape[1])
    self.width.assign(shape[2])
    self.n_samples.assign_add(n_samples)
    self.channel_mse.assign_add(channel_mse)