Oh and amazingly enough, I took the exact same Python file that I built on GPU on my workstation, ran it in Google Collab with a TPU connected, and I really only had to change the configuration to make it work. There was no need to change or tune anything. I was pleasantly surprised by how well that worked.
And let me rant a bit longer, when we started this project I actually first tried to use AWS because our company already had billing set up with them. But AWS was using CUDA 10.0 for SageMaker, which has a known crash- and nan-inducing bug for 2D convolution. AWS support being non-existent as always, we couldn't make things work on AWS and then went bare metal for our GPU training with hetzner.com - which simultaneously also shaved like $30,000 monthly off our bill. AWS's EC2 GPU instances are really expensive when compared to a dedicated root server.