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

## 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

- [x] `torch.jit.script(CrossAttentionBlock(...))` succeeds after fix
- [x] `torch.jit.script(SABlock(...))` succeeds after fix
- [x] `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`

---------

Signed-off-by: Oleksandr Sanin <alexaaander.sanin@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Eric Kerfoot <17726042+ericspod@users.noreply.github.com>

Author: 3 people

monai/networks/blocks/crossattention.py

@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
 
@@ -139,7 +141,7 @@ def __init__(
         )
         self.input_size = input_size
 
-    def forward(self, x: torch.Tensor, context: torch.Tensor | None = None):
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None):  # noqa: UP045
         """
         Args:
             x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C

monai/networks/blocks/selfattention.py

@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -158,7 +160,7 @@ def __init__(
         )
         self.input_size = input_size
 
-    def forward(self, x, attn_mask: torch.Tensor | None = None):
+    def forward(self, x, attn_mask: Optional[torch.Tensor] = None):  # noqa: UP045
         """
         Args:
             x (torch.Tensor): input tensor. B x (s_dim_1 * ... * s_dim_n) x C

monai/networks/blocks/transformerblock.py

@@ -11,6 +11,8 @@
 
 from __future__ import annotations
 
+from typing import Optional
+
 import torch
 import torch.nn as nn
 
@@ -23,6 +25,11 @@ class TransformerBlock(nn.Module):
     An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale <https://arxiv.org/abs/2010.11929>"
     """
 
+    # Treat ``with_cross_attention`` as a TorchScript constant so the cross-attention branch in
+    # ``forward`` is pruned when it is False. Otherwise scripting tries to compile the
+    # ``self.cross_attn(..., context=context)`` call against ``nn.Identity`` and fails.
+    __constants__ = ["with_cross_attention"]
+
     def __init__(
         self,
         hidden_size: int,
@@ -102,7 +109,10 @@ def __init__(
             self.cross_attn = nn.Identity()
 
     def forward(
-        self, x: torch.Tensor, context: torch.Tensor | None = None, attn_mask: torch.Tensor | None = None
+        self,
+        x: torch.Tensor,
+        context: Optional[torch.Tensor] = None,  # noqa: UP045
+        attn_mask: Optional[torch.Tensor] = None,  # noqa: UP045
     ) -> torch.Tensor:
         x = x + self.attn(self.norm1(x), attn_mask=attn_mask)
         if self.with_cross_attention: