Make ViT and Unetr to be torchscript comaptible

[KumoLiu]

Fixes #7936

Description

• Pre-define self.causal_mask = torch.Tensor() before register buffer
• Move norm_cross_attn and cross_attn out of if block

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]

Hi @ericspod and @virginiafdez, the issue is introduced by merging the gen-ai-dev branch: #7886. ViT and Unetr can't support torchscript converting after merging the PR.
Could you please help review this PR? Thanks!

[virginiafdez]

I have checked this with Eric. The only concern we have is in selfattention.py, when you add the CrossAttention block as an attribute by default, it might lead to incompatibilities when loading the module weights from old models, as weights are expected to be there unless you choose strict=False

[KumoLiu]

/build

[ericspod]

Let's put this through now and we'll come back to the issue of weights later, if we're lucky no one's saved models are broken by this anyway.

[KumoLiu]

I have checked this with Eric. The only concern we have is in selfattention.py, when you add the CrossAttention block as an attribute by default, it might lead to incompatibilities when loading the module weights from old models, as weights are expected to be there unless you choose strict=False

Yes, it will potential be an issue, may need to add a load_old_state_dict as @ericspod suggested to fix the loading issue.