Converting tf1 to tf2 for Google ML Course's Recommendation Systems, Exploding Loss

Question

I want to convert this notebook written in tf1 to tf2, more specifically the softmax reccommender model in here:

This is from the ML Course, previous steps descr