What does this error mean and how do I get around it?

I have replicated an architecture from one of the research papers and running it on GPU gives Out Of Memory even on colab (The model is quite deep and huge). So naturally, I want to train it using a TPU.

The same code using GPU doesn’t cause any issue. However it throws this error if I train on TPU:

InvalidArgumentError: 9 root error(s) found. (0) INVALID_ARGUMENT: {{function_node __inference_train_function_56959}} Reshape’s input dynamic dimension is decomposed into multiple output dynamic dimensions, but the constraint is ambiguous and XLA can’t infer the output dimension %reshape.8395 = f32[3,3,2,86,86,32]{5,4,3,2,1,0} reshape(f32[<=18,86,86,32]{3,2,1,0} %convolution.8393), metadata={op_type=”BatchToSpaceND” op_name=”model/conv2d_3/Conv2D/BatchToSpaceND”}.

[[{{node TPUReplicate/_compile/_11021259981217135469/_4}}]]

The colab notebook can be found here:

https://colab.research.google.com/drive/1D-laydrWwLnqAehVREhSkTPoqNXyRGS8?usp=sharing