🐛[BUG]: `EDMPrecondSRV2` called with wrong arguments order

[stathius]

Version

0.7.0a

On which installation method(s) does this occur?

Docker, Source

Describe the issue

Analysis: Dimensions of model and/or data mismatch.

This only affects edmv2

Minimum reproducible example

python3 train.py --config-name=config_train_diffusion.yaml
config: 
     arch: ddpmpp-cwb
     precond: edmv2
     task: diffusion

Relevant log output

**Stack trace:**


Traceback (most recent call last):
  File "/code/modulus/examples/generative/corrdiff/train.py", line 344, in main
    training_loop.training_loop(
  File "/code/modulus/examples/generative/corrdiff/training/training_loop.py", line 267, in training_loop   <--------------
    loss = loss_fn(
  File "/code/modulus/modulus/metrics/diffusion/loss.py", line 336, in __call__
    D_yn = net(y + n, y_lr, sigma, labels, augment_labels=augment_labels)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/nvtx/nvtx.py", line 116, in inner
    result = func(*args, **kwargs)
  File "/code/modulus/modulus/models/diffusion/preconditioning.py", line 891, in forward
    arg = torch.cat([c_in * x, condition], dim=1)
RuntimeError: The size of tensor a (4816896) must match the size of tensor b (2) at non-singleton dimension 0


**Analysis**: Dimensions of model and/or data mismatch.

Environment details

python version: 3.10
modulus commit: `c07fa25321c48a1d71efca12b67d056adbca8bd4`