Tested a run of the notebooks (locally on mac using 'mps')
Very small stuff like adding a logger to import around for logging and random line changes so that can use 'mps' easier and for autodetecting that
A bit of refactoring by moving different components (the activation function and attention mechanism) to their own files which are imported into the model