def METHOD_NAME(cls: T, basic_unit_tag: bool = True) -> T:
    """
    To wrap a module as a basic unit, is to make it a primitive and stop the engine from digging deeper into it.
    ``basic_unit_tag`` is true by default. If set to false, it will not be explicitly mark as a basic unit, and
    graph parser will continue to parse. Currently, this is to handle a special case in ``nn.Sequential``.
    Although ``basic_unit`` calls ``trace`` in its implementation, it is not for serialization. Rather, it is meant
    to capture the initialization arguments for mutation. Also, graph execution engine will stop digging into the inner
    modules when it reaches a module that is decorated with ``basic_unit``.
    .. code-block:: python
        @basic_unit
        class PrimitiveOp(nn.Module):
            ...
    """
    # Internal flag. See nni.trace
    nni_trace_flag = os.environ.get('NNI_TRACE_FLAG', '')
    if nni_trace_flag.lower() == 'disable':
        return cls
    if _check_wrapped(cls, 'basic_unit'):
        return cls
    import torch.nn as nn
    assert issubclass(cls, nn.Module), 'When using @basic_unit, the class must be a subclass of nn.Module.'  # type: ignore
    cls = trace(cls)
    cls._nni_basic_unit = basic_unit_tag  # type: ignore
    _torchscript_patch(cls)
    return cls