Custom loss and visualization

Hi All!
I may be missing something very simple, but have been trying for a while now and not getting anywhere.
I have a custom loss function where I compute a loss based on boundary normals (normal_x, normal_y) and some derivaties of ‘u’ (using tf.gradients etc). - all very similar to tutorial 9 of the Manual.
Inside the custom_loss I have access to x, y of the domain together with normal_x, normal_y through domain_invar.

Now for debugging purposes I would like to be able to plot normal_x, normal_y and my computed derivatives of u in the InferenceDomain so it gets saved to the vtu and I can visualize it.

Whatever I try I always get only ‘x’,‘y’,‘z’,‘area’,‘sdf’ and ‘u’ saved in the (inference) vtu. How can I:

  1. compute some additional variables in the InferenceDomain (i.e. derivatives of ‘u’) and then
  2. save those variables to vtu file

Thanks for any pointers,
Jarek