1478 Fix TorchScript issue in AHnet

[Nic-Ma]

Fixes #1478 .

Description

This PR fixed the TorchScript issue in AHNet when num_block_psp=0.

Status

Ready

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 --codeformat --coverage.
• Quick tests passed locally by running ./runtests.sh --quick.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[Nic-Ma]

/black

[wyli]

thanks, could you add a test case to cover this change?

[Nic-Ma]

/black

[Nic-Ma]

Hi @wyli ,

Thanks for your review.
I updated the PR according to your comments.
Could you please help review it again?

Thanks.