@ptrblck i have reshaped the array to 28*28 before applying the augmentation and augmentation is working fine now. but the interesting thing is. after applying the augmentation model is performing worse, the loss is not at all decreasing. I tried the same augmentation parameters with albumentation and the model does really well in that. even without augmentation also model does well. check the same link for the code.