3432 make vit support torchscript

[yiheng-wang-nv]

Signed-off-by: Yiheng Wang vennw@nvidia.com

Fixes #3432 .

Description

This PR is used to fix #3422 , different from the way used in #3433 which requires to remove einops, this PR keeps this library.

During this PR, the direct calling of einops.arrange in forward functions of class SABlock has been replaced by using Rearrange layers in order to support torchscript.

Status

Ready/Work in progress/Hold

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.

[wyli]

thanks this looks good me, if @ahatamiz has no further comments we merge it today.

[wyli]

some typing issue here https://github.com/Project-MONAI/MONAI/runs/5122538609?check_suite_focus=true#step:7:21211

[ericspod]

There may be some deep-seated incompatibilities between Torchscript and Einops: arogozhnikov/einops#115 The issue with build-gpu/GPU-quick-py3 (PT18+CUDA102) failing may be related, later versions of Pytorch after 1.8 may work, I ran tests with Pytorch 1.10 just now and that version at least works. We may have to state only a certain Pytorch version is scriptable and put the condition on the script test for ViT.

[wyli]

/build