NotImplementedError: Support for training_epoch_end

[Wadha-Almattar]

Hi there, I was trying to reproduce this work. But it keeps giving me this error. Even though, I have changed the method by adding adding [on] to be differ than Lightly implementation. Still, I'm getting the same error.

Ay suggestion will be appreciated,

The error:
NotImplementedError: Support for training_epoch_end has been removed in v2.0.0. MocoModel implements this method. You can use the on_train_epoch_end hook instead.
To access outputs, save them in-memory as instance attributes.
You can find migration examples in Lightning-AI/pytorch-lightning#16520.