def METHOD_NAME(cls, old_forward):
    @functools.wraps(old_forward)
    def new_forward(self, *args, **kwargs):
        if not self.define_subroutine:
            return old_forward(self, *args, **kwargs)
        block_builder = relax.BlockBuilder.current()
        assert block_builder is not None, (
            f"Class {type(self)} has cls.define_subroutines = True, "
            "but is called outsdie of a block_builder environment.  "
            "relax.BlockBuilder.current() is required "
            "to determine where to generate the subroutine."
        )
        func_args = self._normalize_subroutine_args(block_builder, *args, **kwargs)
        subroutine, is_nn_tensor_output = self._get_subroutine(
            block_builder, old_forward, func_args
        )
        subroutine_args = [
            arg._expr if isinstance(arg, nn.Tensor) else arg
            for arg in [*func_args.values(), *self.parameters()]
        ]
        out = subroutine(*subroutine_args)
        if is_nn_tensor_output:
            if out.struct_info_ is None:
                out = block_builder.emit(out, name_hint=f"{subroutine.name_hint}_output")
            out = nn.Tensor(_expr=out)
        return out
    new_forward._is_subroutine_mixin = True
    return new_forward