DynUnet get_feature_maps() issue when using distributed training

Describe the bug If we use DistributedDataParallel to wrap the network, calling the ‘get_feature_maps’ function will raise the following error:

It is common to return multiple variables in a network’s forward function, not only for this case, but also for other situations such as for calculating triplet loss in metric learning based tasks. Therefore, here I think we should re-consider about what to return in the forward function of DynUnet.

Let me think about it and submit a PR then we can discuss it later. @Nic-Ma @wyli @rijobro

[yiheng-wang-nv]

Hi @rijobro , thanks for the advice, but return this kind of dict doesn’t fix the torchscript issue as @Nic-Ma mentioned. I suggest here we still return a list, and add the corresponding docstrings. In order to help users use this network, I will also add/update tutorials. Let me submit a PR first for you to review.

[Nic-Ma]

Hi @rijobro ,

I think most of the previous problems are due to the list output during validation or inference. So I suggest to return a list of data during training, return only the output instead of [out] during validation or inference.

Thanks.