1111 Fix TorchScript load issue in DenseNet

[Nic-Ma]

Fixes #1111 .

Description

This PR fixed the TorchScript load key issue in the DenseNet.

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.
• In-line docstrings updated.
• Documentation updated, tested make html command in the docs/ folder.

[ericspod]

Perhaps we should include unit tests to verify our models are Torchscript compatible at least in their default forms. A test function to script and reload a network then test its forward pass:

from io import BytesIO


def test_script_save(net, *inputs):
    scripted = torch.jit.script(net)
    buffer = scripted.save_to_buffer()
    reloaded_net = torch.jit.load(BytesIO(buffer))
    
    return net(*inputs), reloaded_net(*inputs)


t = torch.rand(1,5)
net = torch.nn.Linear(5,6)

out_orig, out_reloaded = test_script_save(net, t)
assert torch.allclose(out_orig, out_reloaded)

[Nic-Ma]

Hi @ericspod ,

Thanks for your review, very good idea.
I submitted a ticket to track this task #1117, and I will submit another separate PR next week to check TorchScript in unit tests of all the networks.

Thanks.