Training AlexNet like structure on entire ImageNet in Tensorflow 2.0, Problem with flat loss

Question

I am trying to replicate AlexNet model on ImageNet (for learning purposes). The dataset is 1.2 million ImageNet dataset with 50K validation. This is my code :

model funct