LightningModule docs is wrong when defining "Training with DataParallel"
The docs says when training with DataParallel the training_step_end function should be like this:
def training_step_end(self, batch_parts):
    gpu_0_prediction = batch_parts[0]['pred']
    gpu_1_prediction = batch_parts[1]['pred']
    # do something with both outputs
    return (batch_parts[0]['loss'] + batch_parts[1]['loss']) / 2

but acually, when i try it for myself, it should be like this:
def training_step_end(self, batch_parts):
    gpu_0_prediction = batch_parts['pred'][0]
    gpu_1_prediction = batch_parts['pred'][1]
    # do something with both outputs
    return (batch_parts['loss'][0] + batch_parts['loss'][1]) / 2

yes, the order of the axes are opposite.