Well, I haven't "switched" since I've been using Tensorflow since the start of the project. I was just curious to see if Pytorch could allow me to squeeze more juice and after spending a weekend trying to learn assembly Pytorch syntax it turns out that yes, but actually no. So yeah I'm perfectly content with using model.fit and calling it a day for the time being.

Oh and I also forgot: Pytorch won't train with a distributed strategy in a Jupyter environment. KEK.