How to write a squared loss function for categorisation in Keras/Tensorflow?

Question

I have one hot encoded data and I want a loss function that calculates (1 - predicted probability for the correct category) ^ 2 for each data point, my simple attempt was