Tensorflow Training Crashes in last step of first epoch for audio classifier

Question

I was trying to convert a custom DNN network originally written in Pytorch to Tensorflow 2 that has input shape required as (batchsize,39,101,1). After using the feature e