Error while training on colab TPU

[waleedrazakhan92]

While trying to train I keep getting this error
File "/content/stylegan2-tpu/dnnlib/tflib/autosummary.py", line 116, in image_grid (grid_shape, int(input_tensor.shape[0]))) ValueError: Grid shape (3, 3) incompatible with minibatch size 8.

First I tried with a larger batch size but kept getting File "/tensorflow-1.15.2/python3.7/tensorflow_core/python/framework/ops.py", line 1748, in __init__ self._traceback = tf_stack.extract_stack() this error. So I googled around and got to know that its an overflow error. So i set the BATCH_PER to 1 and tried training again but now I'm getting the incompatible minibatch size error mentioned above.
Please let me know if I'm doing anything wrong. Or the grid size to save is somehow hardcoded and not working with the new BATCH_PER size. I'm using image size to be 256x256