def METHOD_NAME(self):
    with tf.device("/device:GPU:0"):
        mean_obj = Mean()
        y_true_0 = tf.constant([0, 1, 2], dtype=tf.float32)
        y_true_1 = tf.constant([3, 4], dtype=tf.float32)
        self.evaluate(
            mean_obj.update_state(
                y_true=y_true_0, y_pred=tf.constant(0, dtype=tf.float32)
            )
        )
        self.evaluate(
            mean_obj.update_state(
                y_true=y_true_1, y_pred=tf.constant(0, dtype=tf.float32)
            )
        )
    self.assertEqual(2, self.evaluate(mean_obj.result()))
    if not tf.executing_eagerly():
        self.assertEndsWith(y_true_0.device, "/device:GPU:0")
        self.assertEndsWith(y_true_1.device, "/device:GPU:0")
    self.assertEndsWith(mean_obj.values[0].device, "/device:CPU:0")
    self.assertEndsWith(mean_obj.values[1].device, "/device:CPU:0")