From 9e3cf855925c57a2747d73be0c362a3f4fc9939b Mon Sep 17 00:00:00 2001 From: Nic Ma Date: Sat, 17 Oct 2020 19:32:49 +0800 Subject: [PATCH] [DLMED] fix torchscript issue in densenet Signed-off-by: Nic Ma --- monai/networks/nets/densenet.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/monai/networks/nets/densenet.py b/monai/networks/nets/densenet.py index 2ec6efcb44..0d1aa8c447 100644 --- a/monai/networks/nets/densenet.py +++ b/monai/networks/nets/densenet.py @@ -174,7 +174,7 @@ def __init__( ("relu", nn.ReLU(inplace=True)), ("norm", avg_pool_type(1)), ("flatten", nn.Flatten(1)), - ("class", nn.Linear(in_channels, out_channels)), + ("out", nn.Linear(in_channels, out_channels)), ] ) )