training_epoch_end error

[jimmylihui]

When I run the code, the following error appear:
NotImplementedError: Support for training_epoch_end has been removed in v2.0.0. TSDiff 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.

[abdulfatir]

@jimmylihui Did you install the package as in the instructions? Do you encounter this issue with version pinned in https://github.com/amazon-science/unconditional-time-series-diffusion/blob/main/pyproject.toml?