In my projects, and in kaggle competitions, pytorch code is always the easiest and shortest part (yet obviously really complicated theoretically).

There's no reason to switch to another framework like lightning that will disappear in a couple of years.