Batchsize and learning rate scheduler
I was wondering if I need to adjust the batch size when using TPUs. I had a memory error when trying to run a (image 256x256x3) batch of size 128, which works perfectly fine on GPUs.
Furthermore, do I need to adjust my custom learning rate scheduler which on GPUs run every batch (not just epoch):
def configure_optimizers(self):
        optimizer = get_optimizer(self)
        one_cycle_scheduler = OneCycleScheduler(optimizer, EPOCHS * steps_per_epoch // GRAD_ACCUMULATE)
        scheduler = {'scheduler': one_cycle_scheduler, "interval": "step"}

        return [optimizer], [scheduler]

pl version: 0.8.5
system: kaggle