Fix example on U-Net learning

[chaithyagr]

While the example is good, some math rendering and some images arent rendered. We did see it in local tests to be fine. Perhaps we missed something,