How to use EMA callback for pytorch-lighting?

[KasuganoLove]

The EMA callback is great! I want use it for training and test for pytorch-lightining. My old code is like blow, how should I change it?

Train:

# train the model
def load_callbacks(args):
    callbacks = []

    callbacks.append(plc.ModelCheckpoint(
        monitor='val_PSNR',
        filename='best-{epoch:02d}-{val_PSNR:.3f}',
        mode='max',
        save_last=True,
        save_on_train_epoch_end=True,
        every_n_epochs=50
    ))

    callbacks.append(plc.LearningRateMonitor(logging_interval='epoch'))

    return callbacks

trainer = Trainer(
    log_every_n_steps=1,
    accelerator='gpu',
    devices=2,
    max_epochs=5000,
    check_val_every_n_epoch=50,
    callbacks=load_callbacks(args),
    deterministic='warn',
    strategy='ddp'
)
trainer.fit(model=my_model, datamodule=data_module)

Test:

# test the model
my_model.load_from_checkpoint(checkpoint_path=args.model_ckpt)

trainer = Trainer(
    devices=1,
    logger=True,
)

trainer.validate(model=my_model, datamodule=data_module)
trainer.test(model=my_model, datamodule=data_module)