def test_summary_training_mode():
    """Test that the model summary captures the training mode on all submodules."""
    model = DeepNestedModel()
    model.branch1[1][0].eval()
    model.branch2.eval()

    summary = summarize(model, max_depth=1)
    summary_data = OrderedDict(summary._get_summary_data())
    assert summary_data["Mode"] == [
        "train",  # branch1
        "eval",  # branch2
        "train",  # head
    ]
    assert summary.total_training_modes == {"train": 17, "eval": 2}

    summary = summarize(model, max_depth=-1)
    expected_eval = {"branch1.1.0", "branch2"}
    for name, layer_summary in summary._layer_summary.items():
        assert (name in expected_eval) == (not layer_summary.training)

    # A model with params not belonging to a layer
    model = NonLayerParamsModel()
    model.layer.eval()
    summary = summarize(model)
    summary_data = OrderedDict(summary._get_summary_data())
    assert summary_data["Mode"] == ["eval", "n/a"]
    assert summary.total_training_modes == {"train": 0, "eval": 1}
Test that the model summary captures the training mode on all submodules.
tests/tests_pytorch/utilities/test_model_summary.py