Remove metrics hack for dist strat

official/transformer/v2/metrics.py

@@ -160,10 +160,6 @@ class MetricLayer(tf.keras.layers.Layer):
 
   def call(self, inputs):
     logits, targets = inputs[0], inputs[1]
-    # TODO(guptapriya): Remove this check when underlying issue to create
-    # metrics with dist strat in cross replica context is fixed.
-    if (tf.distribute.has_strategy() and
-        not tf.distribute.in_cross_replica_context()):
     for mean, fn in self.metric_mean_fns:
       m = mean(*fn(logits, targets))
       self.add_metric(m)