C2W3 Project Assignment - Dead kernel

While trying to run the train function that calls the train_one_step, my kernel just dies after a few steps. And it does not allow my train function to be subsequently graded by the grader.

I also do get this warning before the training code runs -

WARNING:tensorflow:Gradients do not exist for variables [‘bn_conv1/moving_mean:0’, ‘bn_conv1/moving_variance:0’, ‘bn2a_branch2a/moving_mean:0’, ‘bn2a_branch2a/moving_variance:0’, ‘bn2a_branch2b/moving_mean:0’,

I ran the assignment and I dont get any errors or warnings. Given what the error says I would think there is train_one_step() function espeacially with the gradient calculation…