How do you use the custom generator to train the TensorFlow model on PageMaker?

[dsnsabari]

I am trying to train a TensorFlow custom classification model using a custom data generator on SageMaker Studio. However, the notebook did not utilize the GPU to train the model.

My instance type:

I have stored the images in a zip file on an S3 bucket and unzipped them during training, using the temporary path in the custom generator for training.

I want to add data augmentation while training the model.

The example below loads the images in numpy format and passes them directly to the training model. I want to load the images from the S3 bucket in zip format and train with a custom generator and augmentation.

Link to the example