Fix bug for crashing when positional embedding disabled SongUNet

Open daviddpruitt opened this issue 1 year ago • 1 comments

Modulus Pull Request

Description

SongUNetPosEmbed has the option to disable positional embeding by setting N_grid_channels to 0. This change keeps the forward pass from crashing by adding a check to make sure the positional embedding exists. It also elimates some code duplication between the positional embed and non-positional embed version.

Closes #504 https://github.com/NVIDIA/modulus/issues/504

Checklist

[x] I am familiar with the Contributing Guidelines.
[x] New or existing tests cover these changes.
[x] The documentation is up to date with these changes.
[ ] The CHANGELOG.md is up to date with these changes.
[x] An issue is linked to this pull request.

Dependencies

May 21 '24 00:05 daviddpruitt

/blossom-ci

May 21 '24 00:05 mnabian

@akshaysubr we do have unit tests for this model: https://github.com/NVIDIA/modulus/blob/main/test/models/diffusion/test_song_unet_pos_embd.py#L41

May 21 '24 21:05 mnabian

/blossom-ci

May 30 '24 00:05 daviddpruitt