I just read over this whole thread and noticed that we didn’t really conclusively explain why it doesn’t work with tf.square but it does work with **2. There is one post on this thread that actually does describe the error here, but it’s worth also pointing to this other thread which gives a demonstration of the difference.
1 Like