def METHOD_NAME(self, num_replicas=1):
    config = super().METHOD_NAME(num_replicas=num_replicas)
    return config_util.merge_config(
        config,
        {
            "infer": {
                "length_bucket_width": 1  # To ensure fixed length in each batch.
            }
        },
    )