fix R_Square shape issue in model.evaluate (#2034)

* fix R_Square shape issue in model.evaluate

* add test

* remove erroneous line in test

Author: Jake Schmidt

tensorflow_addons/metrics/r_square.py

@@ -139,4 +139,4 @@ def result(self) -> tf.Tensor:
 
     def reset_states(self) -> None:
         # The state of the metric will be reset at the start of each epoch.
-        K.batch_set_value([(v, 0) for v in self.variables])
+        K.batch_set_value([(v, tf.zeros_like(v)) for v in self.variables])

tensorflow_addons/metrics/tests/r_square_test.py

@@ -117,3 +117,13 @@ def test_r2_sklearn_comparison():
 def test_unrecognized_multioutput():
     with pytest.raises(ValueError):
         initialize_vars(multioutput="meadian")
+
+
+def test_keras_fit():
+    model = tf.keras.Sequential([tf.keras.layers.Dense(1)])
+    model.compile(loss="mse", metrics=[RSquare(y_shape=(1,))])
+    data = tf.data.Dataset.from_tensor_slices(
+        (tf.random.normal(shape=(100, 1)), tf.random.normal(shape=(100, 1)))
+    )
+    data = data.batch(10)
+    model.fit(x=data, validation_data=data)