import tensorflow as tf

                    int(maximum_lower_multiple(x // num_replicas_in_sync, 8) * num_replicas_in_sync)
                    for x in bucket_batch_sizes]
            else:
                bucket_batch_sizes = int(maximum_lower_multiple(
                    bucket_batch_sizes // num_replicas_in_sync, 8) * num_replicas_in_sync)
            return dataset_utils.batch_examples_by_token(
                dataset,
                bucket_boundaries=bucket_boundaries,
                bucket_batch_sizes=bucket_batch_sizes,
                padding_values=padding_values,
                example_length_func=lambda x: {"feature": tf.size(x["feature"]),
                                               "label": tf.size(x["label"])},
                extra_padded_shapes={"src_lang": [], "trg_lang": []}
            )

    def build_metric_layer(self):
        return [SequenceTokenMetricLayer("src"), SequenceTokenMetricLayer("trg"),
                BatchCountMetricLayer("src")]

    def get_eval_metric(self, args, name="metric", ds=None):
        """ Returns a neurst.metrics.metric.Metric object for evaluation."""
        if ds is None or not hasattr(ds, "trg_lang") or ds.trg_lang is None:
