Only have contigous calls after attention blocks

[marksgraham]

Towards #7227 .

Description

There were lots of contigous calls in the DiffusionModelUnet. It turns out these are necessary after attention blocks, as the einops operation sometimes leads to non-contigous tensors that can cause errors. I've tidied the code up so the .contiguous calls are only after attention calls.

A few sentences describing the changes proposed in this pull request.

Types of changes

• Non-breaking change (fix or new feature that would not break existing functionality).
• Breaking change (fix or new feature that would cause existing functionality to change).
• New tests added to cover the changes.
• Integration tests passed locally by running ./runtests.sh -f -u --net --coverage.
• Quick tests passed locally by running ./runtests.sh --quick --unittests --disttests.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[KumoLiu]

/build

[KumoLiu]

One blossom error in TestDiffusionModelUNet3D cc @marksgraham

[2024-05-14T06:55:49.127Z] ======================================================================
[2024-05-14T06:55:49.127Z] ERROR: test_shape_unconditioned_models_4 (tests.test_spade_diffusion_model_unet.TestDiffusionModelUNet3D)
[2024-05-14T06:55:49.127Z] ----------------------------------------------------------------------
[2024-05-14T06:55:49.127Z] Traceback (most recent call last):
[2024-05-14T06:55:49.127Z]   File "/usr/local/lib/python3.8/dist-packages/parameterized/parameterized.py", line 620, in standalone_func
[2024-05-14T06:55:49.127Z]     return func(*(a + p.args), **p.kwargs, **kw)
[2024-05-14T06:55:49.127Z]   File "/home/jenkins/agent/workspace/MONAI-premerge/monai/tests/test_spade_diffusion_model_unet.py", line 519, in test_shape_unconditioned_models
[2024-05-14T06:55:49.127Z]     result = net.forward(
[2024-05-14T06:55:49.127Z]   File "/home/jenkins/agent/workspace/MONAI-premerge/monai/monai/networks/nets/spade_diffusion_model_unet.py", line 932, in forward
[2024-05-14T06:55:49.127Z]     h = upsample_block(hidden_states=h, res_hidden_states_list=res_samples, seg=seg, temb=emb, context=context)
[2024-05-14T06:55:49.127Z]   File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
[2024-05-14T06:55:49.127Z]     return forward_call(*input, **kwargs)
[2024-05-14T06:55:49.127Z]   File "/home/jenkins/agent/workspace/MONAI-premerge/monai/monai/networks/nets/spade_diffusion_model_unet.py", line 436, in forward
[2024-05-14T06:55:49.127Z]     hidden_states = self.upsampler(hidden_states, temb)
[2024-05-14T06:55:49.127Z]   File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
[2024-05-14T06:55:49.127Z]     return forward_call(*input, **kwargs)
[2024-05-14T06:55:49.127Z]   File "/home/jenkins/agent/workspace/MONAI-premerge/monai/monai/networks/nets/diffusion_model_unet.py", line 396, in forward
[2024-05-14T06:55:49.127Z]     h = self.norm1(h)
[2024-05-14T06:55:49.127Z]   File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
[2024-05-14T06:55:49.127Z]     return forward_call(*input, **kwargs)
[2024-05-14T06:55:49.127Z]   File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/normalization.py", line 268, in forward
[2024-05-14T06:55:49.127Z]     return F.group_norm(
[2024-05-14T06:55:49.127Z]   File "/usr/local/lib/python3.8/dist-packages/torch/nn/functional.py", line 2360, in group_norm
[2024-05-14T06:55:49.127Z]     return torch.group_norm(input, num_groups, weight, bias, eps, torch.backends.cudnn.enabled)
[2024-05-14T06:55:49.127Z] RuntimeError: Unsupported memory format. Supports only ChannelsLast, Contiguous
[2024-05-14T06:55:49.127Z] 
[2024-05-14T06:55:49.127Z] ----------------------------------------------------------------------

[marksgraham]

@KumoLiu should be fixed now

[KumoLiu]

/build