Typo in sample codes of training_step_end() documentation
📚 Documentation
def training_step(self, batch, batch_idx):
    x, y = batch
    y_hat = self.model(x)
    loss = F.cross_entropy(y_hat, y)
    pred = ...
    return {'loss': loss, 'pred': pred}

def training_step_end(self, batch_parts):
    gpu_0_prediction = batch_parts.pred[0]['pred'] # unexpected `pred` here
    gpu_1_prediction = batch_parts.pred[1]['pred'] # unexpected `pred` here

    # do something with both outputs
    return (batch_parts[0]['loss'] + batch_parts[1]['loss']) / 2
It should be to avoid misleading.
def training_step_end(self, batch_parts):
    gpu_0_prediction = batch_parts[0]['pred'] # unexpected `pred` here
    gpu_1_prediction = batch_parts[1]['pred'] # unexpected `pred` here

    # do something with both outputs
    return (batch_parts[0]['loss'] + batch_parts[1]['loss']) / 2