'Generate predictions/logits for all shards.
Classes with cross-shard interaction will override this function.
Args:
sharded_body_output: A list of Tensors.
sharded_targets: A list of Tensors.
data_parallelism: a expert_utils.Parallelism object.
Returns:
sharded_logits: A list of Tensors.'

def top_sharded(self, sharded_body_output, sharded_targets, data_parallelism):

    return data_parallelism(self.top, sharded_body_output, sharded_targets)
