Minimal examples for PKM layer training contains an implementation of PKM layer training for both pytorch and tensorflow2