fix(networks): replace Tensor | None with Optional[Tensor] for TorchScript compatibility

[AlexanderSanin]

Summary

Fixes #7939

The | union type syntax (e.g. torch.Tensor | None) was introduced in Python 3.10. While from __future__ import annotations defers annotation evaluation at runtime, TorchScript's annotation parser does not support this syntax and fails when scripting models that include these forward method signatures, producing:

RuntimeError: Can't redefine method: forward on class

This PR replaces torch.Tensor | None with Optional[torch.Tensor] (from typing) in the forward methods of the three blocks that form the ViT/UNETR scripting path:

• monai/networks/blocks/crossattention.py — CrossAttentionBlock.forward
• monai/networks/blocks/selfattention.py — SABlock.forward
• monai/networks/blocks/transformerblock.py — TransformerBlock.forward

Test plan

• torch.jit.script(CrossAttentionBlock(...)) succeeds after fix
• torch.jit.script(SABlock(...)) succeeds after fix
• torch.jit.script(TransformerBlock(...)) succeeds after fix
• python -m pytest tests/networks/nets/test_unetr.py -k test_script
• python -m pytest tests/networks/blocks/test_crossattention.py
• python -m pytest tests/networks/blocks/test_selfattention.py

[AlexanderSanin]

Hey @holgerroth @wyli @ericspod. Could you, please, have a look at this?

[coderabbitai]

Caution

Review failed

The head commit changed during the review from eb55abe to a80169b.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[ericspod]

Hi @AlexanderSanin I don't see files changed anymore since pre-commit (using ruff I think) autofixed your changes back to what they were. You may have to create a variable for your type with OptTensor: TypeAlias = Optional[torch.Tensor] at the tope of the source files with an annotation to make ruff ignore it. This can then be used where you want.

However, we don't support Python 3.9 anymore so we expect the X | None syntax to always work, and Torchscript itself is deprecated and will be phased out of MONAI soon. If you desperately need Torchscript we can look at this fix but I'd suggest migrating off of it instead.