Trained weights are on CPU despite the model being trained on GPU

LightningModule
3

Absolutely, here are the trainer settings:
trainer = pl.Trainer(max_epochs=EPOCHS,log_every_n_steps=10, profiler="simple", logger=logger,accelerator="gpu", devices=1)