def test_pytorch_profiler_nested(tmp_path):
    """Ensure that the profiler handles nested context."""
    kwargs = {} if _TORCH_GREATER_EQUAL_2_4 else {"use_cuda": False}
    pytorch_profiler = PyTorchProfiler(dirpath=tmp_path, filename="profiler", schedule=None, **kwargs)

    with pytorch_profiler.profile("a"):
        a = torch.ones(42)
        with pytorch_profiler.profile("b"):
            b = torch.zeros(42)
        with pytorch_profiler.profile("c"):
            _ = a + b

    pytorch_profiler.describe()

    events_name = {e.name for e in pytorch_profiler.function_events}

    names = {"[pl][profile]a", "[pl][profile]b", "[pl][profile]c"}
    ops = {"add", "empty", "fill_", "ones", "zero_", "zeros"}
    ops = {"aten::" + op for op in ops}

    expected = names.union(ops)
    assert events_name == expected, (events_name, torch.__version__, platform.system())
Ensure that the profiler handles nested context.
tests/tests_pytorch/profilers/test_profiler.py