v0.5.2: Fix dependency

Pytorch-lightning depreciated a feature which broke training; this is now fixed.

• prelim TPU support + more correct pytorch-lightning training (#105)
• Bump min pytorch-lightning version to 1.3.1
• fixes for schema generation