Nakul,

Please mention the TF version that you are using. I had used TF 1.14 for the code

Thanks,

Renu