Thank you for sharing this wonderful works!

I found a simple lightning dependancy and epoch parameter in the _cosine_scheduler_fn(...) in Marlin>model>marlin.py cause me an error.

Somehow the epoch parameter above called to a number of max_epochs not max_epochs-1. (This is not a lightning's bug or something)

So I added a simple script to _cosine_scheduler_fn(..) like belows:

def _cosine_scheduler_fn(self, epoch):
    if epoch > 0: epoch -= 1 # lightning dependancy
    return self.lr_scheduler_factors[epoch]

Again, many thanks for your great achievement to the deep learning commnuniities!