Switch to pytorch lighting training loop

With Iseue 4 done, we could easily switch to pytorch lightning to make the code structure cleaner and more robust.