Repository: Project-MONAI/MONAI
File: monai/networks/utils.py (Python)

Code before:
    Set network(s) to train mode and then return to original state at the end.

    Args:
        nets: Input network(s)

    Examples

    .. code-block:: python

        t=torch.rand(1,1,16,16)
        p=torch.nn.Conv2d(1,1,3)
        p.eval()
        print(p.training)  # False
        with train_mode(p):
            print(p.training)  # True
            print(p(t).sum().backward())  # No exception
    """

    # Get original state of network(s)
    # Check the training attribute in case it's TensorRT based models which don't have this attribute.
    eval_list = [n for n in nets if hasattr(n, "training") and (not n.training)]

    try:
        # set to train mode
        with torch.set_grad_enabled(True):
            yield [n.train() if hasattr(n, "train") else n for n in nets if hasattr(n, "train")]
    finally:
        # Return required networks to eval_list
        for n in eval_list:
            if hasattr(n, "eval"):
                n.eval()


def get_state_dict(obj: torch.nn.Module | Mapping):
    """
    Get the state dict of input object if has `state_dict`, otherwise, return object directly.
    For data parallel model, automatically convert it to regular model first.

    Args:
        obj: input object to check and get the state_dict.

    """
    if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
        obj = obj.module
    return obj.state_dict() if hasattr(obj, "state_dict") else obj


def copy_model_state(
    dst: torch.nn.Module | Mapping,
    src: torch.nn.Module | Mapping,
    dst_prefix="",

Review comment (suggestion):
```suggestion
            yield [n.train() if hasattr(n, "train") else n for n in nets]
```

Code after:
    Set network(s) to train mode and then return to original state at the end.

    Args:
        nets: Input network(s)

    Examples

    .. code-block:: python

        t=torch.rand(1,1,16,16)
        p=torch.nn.Conv2d(1,1,3)
        p.eval()
        print(p.training)  # False
        with train_mode(p):
            print(p.training)  # True
            print(p(t).sum().backward())  # No exception
    """

    # Get original state of network(s)
    # Check the training attribute in case it's TensorRT based models which don't have this attribute.
    eval_list = [n for n in nets if hasattr(n, "training") and (not n.training)]

    try:
        # set to train mode
        with torch.set_grad_enabled(True):
            yield [n.train() if hasattr(n, "train") else n for n in nets]
    finally:
        # Return required networks to eval_list
        for n in eval_list:
            if hasattr(n, "eval"):
                n.eval()


def get_state_dict(obj: torch.nn.Module | Mapping):
    """
    Get the state dict of input object if has `state_dict`, otherwise, return object directly.
    For data parallel model, automatically convert it to regular model first.

    Args:
        obj: input object to check and get the state_dict.

    """
    if isinstance(obj, (nn.DataParallel, nn.parallel.DistributedDataParallel)):
        obj = obj.module
    return obj.state_dict() if hasattr(obj, "state_dict") else obj


def copy_model_state(
    dst: torch.nn.Module | Mapping,
    src: torch.nn.Module | Mapping,
    dst_prefix="",