def test_total_training_modes():
    """Test that the `total_training_modes` counts the modules in 'train' and 'eval' mode, excluding the root
    module."""

    class ModelWithoutChildren(LightningModule):
        pass

    summary = ModelSummary(ModelWithoutChildren())
    assert summary.total_training_modes == {"train": 0, "eval": 0}

    model = DeepNestedModel()
    summary = ModelSummary(model)
    assert summary.total_training_modes == {"train": 19, "eval": 0}
    assert sum(summary.total_training_modes.values()) == len(list(model.modules())) - 1

    model = DeepNestedModel()
    summary = ModelSummary(model)
    model.branch1[1][0].eval()
    model.branch2.eval()
    assert summary.total_training_modes == {"train": 17, "eval": 2}
    assert sum(summary.total_training_modes.values()) == len(list(model.modules())) - 1
Test that the `total_training_modes` counts the modules in 'train' and 'eval' mode, excluding the root
    module.
tests/tests_pytorch/utilities/test_model_summary.py