How to retrieve DeepONet information, such as Loss and Weights split by branch and trunk?

[Daniel-EST]

I wonder if it is possible to retrieve information about the model on each epoch. I was able to return some information (Weights, Bias and Optimizer) by using Model.state_dict() in a custom Callback. However, I am not sure if the branch and trunk model layers are on a specific order.

Moreover, It seems like I am only able to access the total loss (m.s.e) and not the branch and trunk loss separately.

[lululxvi]

Branch and trunk nets don't have separate losses.