🐛[BUG]: Diffusion network, SongUNetPosEmbed tries to apply positional embedding when its disabled.

[daviddpruitt]

Version

0.6.0

On which installation method(s) does this occur?

Docker, Pip, Source

Describe the issue

When using SongUNetPosEmbed for diffusion there is an option to disable positional embedding by setting N_grid_channels to 0. If the positional embedding is disabled the forward pass will still apply the positional embedding resulting in a crash. See: https://github.com/NVIDIA/modulus/blob/f49ceae9b30df5fa065976d2fcca66e19f98ace1/modulus/models/diffusion/song_unet.py#L516

Minimum reproducible example

No response

Relevant log output

Traceback (most recent call last):
  File "/root/modulus_corr_patch/examples/generative/corrdiff/train.py", line 352, in <module>
    main()
  File "/usr/local/lib/python3.10/dist-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/usr/local/lib/python3.10/dist-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/usr/local/lib/python3.10/dist-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/usr/local/lib/python3.10/dist-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/root/modulus_corr_patch/examples/generative/corrdiff/train.py", line 344, in main
    training_loop.training_loop(
  File "/root/modulus_corr_patch/examples/generative/corrdiff/training/training_loop.py", line 288, in training_loop                                                                                                    loss = loss_fn(
  File "/usr/local/lib/python3.10/dist-packages/modulus/metrics/diffusion/loss.py", line 411, in __call__                                                                                                               D_yn = net(input, y_lr, sigma, labels, augment_labels=augment_labels)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl                                                                                                           return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/modulus/models/diffusion/unet.py", line 152, in forward
    F_x = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl                                                                                                           return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1527, 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 "/usr/local/lib/python3.10/dist-packages/modulus/models/diffusion/song_unet.py", line 516, in forward                                                                                                            selected_pos_embd = self.positional_embedding_indexing(x, global_index)
  File "/usr/local/lib/python3.10/dist-packages/modulus/models/diffusion/song_unet.py", line 578, in positional_embedding_indexing                                                                                      self.pos_embd.to(x.dtype)
AttributeError: 'NoneType' object has no attribute 'to'

Environment details

No response