Currently, I'm using tf.keras.utils.image_dataset_from_directory to load about 280,000+ images with a batch size 128 and for the model, I am using a batch size of 1024.
However, I am getting "E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node AssignAddVariableOp.".
On doing a simple google search,
I am assuming that "The Device" i.e. TPU might not able to handle the I/O operation and the fix should be reducing batch size.
However, I have tried setting the dataset's batch size 32, 64, 128, 256, 1024 and for the model batch size, 1024, 128, 32, but seem to get the same error.So basically using a large batch size and reducing by halving it. I did find this method to calculate, however I am not sure how to calculate the size of tensor from tf.keras.utils.image_dataset_from_directory.