normalization for layer output implementation

[tcsmaster]

Hi all,

I am researching the effects of normalization for layer outputs in neural networks, similar to this paper: https://arxiv.org/pdf/2209.01018.pdf, but applying it to PDEs.
I have been using my custom architecture, but I've run into problems, which would be an easy fix if I was using DeepXDE (namely switching from one optimizer to another).

My question is: do you think it'd be easier to modify the source code to achieve the scaling, or should I stick to my custom implementation, and try to implement the switching?

[lululxvi]

DeepXDE is very easy to modify.