🐛[BUG]: DeepONetConstraint.from_numpy not compatible with IntegralLossNorm #195

Version

1.6.0

On which installation method(s) does this occur?

Pip

Describe the issue

I am trying to use loss=IntegralNormLoss() in the DeepONetConstraint.from_numpy constraint to set an integral continuity plane constraint for DeepONets. However, IntegralLossNorm expects a list of dictionaries for its invar and outvar, whereas DeepONetConstraint.from_numpy must be given dictionaries. I am not sure why there is this inconsistency. Perhaps you are not supposed to use IntegralLossNorm as the loss in the constraint?
It seems that there is a version of IntegralLossNorm._loss that uses dictionaries beneath the current version in the code, so perhaps there is an "if" statement missing or something simple like that. Using this code instead, my training script actually runs.

Minimum reproducible example

int_cont = DeepONetConstraint.from_numpy(
nodes=nodes,
invar=invar,
outvar=outvar,
batch_size=cfg.batch_size.integral_planes,
loss=IntegralLossNorm(2),
)

Relevant log output

Error executing job with overrides: []
Traceback (most recent call last):
File "/home/ahal/GitRepos/MLCFD/pinn/pipe_bend/DeepONet_pipe_bend.py", line 732, in run
flow_slv.solve()
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/solver/solver.py", line 175, in solve
self._train_loop(sigterm_handler)
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/trainer.py", line 545, in _train_loop
loss, losses = self._cuda_graph_training_step(step)
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/trainer.py", line 726, in _cuda_graph_training_step
self.loss_static, self.losses_static = self.compute_gradients(
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/trainer.py", line 78, in adam_compute_gradients
losses_minibatch = self.compute_losses(step)
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/solver/solver.py", line 68, in compute_losses
return self.domain.compute_losses(step)
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/domain/domain.py", line 153, in compute_losses
for loss_key, value in constraint.loss(step).items():
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/domain/constraint/continuous.py", line 139, in loss
losses = self._loss(
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
return forward_call(*args, **kwargs)
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/loss/loss.py", line 172, in forward
return IntegralLossNorm._loss(
File "/home/ahal/env/modulus3.10/lib/python3.10/site-packages/modulus/sym/loss/loss.py", line 141, in _loss
losses = {key: 0 for key in list_pred_outvar[0].keys()}
KeyError: 0

Environment details

No response

Other/Misc.

Writing a minimal code example that reproduces this will require considerable work from my part. I have only provided the constraint I am setting. Hope you can find out the rest for yourself.