training_epoch_end only returning the last train batch

[armoses]

I have a pytorch lightning module that includes the following section:

`def training_step(self, batch, batch_idx):
losses, tensors = self.shared_step(batch)
return losses

def validation_step(self, batch, batch_idx):
    losses, tensors = self.shared_step(batch)
    return losses, tensors

def training_epoch_end(self, losses):
    my_function(losses)

def validation_epoch_end(self, outputs):
    my_function2(outputs)`

The losses variable is a dictionary that includes the key "loss". The input to validation_epoch_end() is a list of (losses, tensors) tuples from each of the batches as expected. The input to training_epoch_end() is a list of the correct length (the number of batches), but every element is the same losses dictionary from the final batch.

How can I input the losses from every train batch into the training_epoch_end() method (like the inputs to validation_epoch_end())?

[ananthsub]

This issue was reported in #8603 - are you able to try Lightning 1.4.1, which contains the fix?

And for broader discussion on these hooks, and alternatives you have to access the per-step outputs, see #8731

[armoses]

Ahh yes it works with 1.4.1! I totally missed that issue. Thank you!