Project running out of memory using code from C3 W1 Lab 2

So I am trying to train a model that detects whether someone is using sunglasses using tfx and a subset of the celebsA dataset (~ 26k images). I have written the images and labels to a TFrecord that is 232.9MB.

When I then go through the different components, using code almost identical to that found in C3 W1 Lab 2, I always run out of memory when running the Transform component. Is this normal? By the way, I am running this on a TPU with 32GB of RAM as I am using Colab Pro.

If so, what would be the best way to overcome the problem? Just create many smaller records and one by one go through the components on them?

Since this is an image dataset, you should be alright with just on transformation i.e. rescaling each pixel by dividing by 255. Why use tfx instead of dataset?