def METHOD_NAME(
    self, model: torch.nn.Module, context_to_use: Optional["TracingContext"] = None, as_eval: bool = False
) -> DynamicGraph:
    sd = deepcopy(model.state_dict())
    from nncf.torch.dynamic_graph.context import TracingContext  # pylint: disable=cyclic-import
    if context_to_use is None:
        context_to_use = TracingContext()
    context_to_use.enable_trace_dynamic_graph()
    from nncf.torch.utils import training_mode_switcher  # pylint: disable=cyclic-import
    with context_to_use as _ctx:
        _ctx.base_module_thread_local_replica = model
        with torch.no_grad():
            if as_eval:
                with training_mode_switcher(model, is_training=False):
                    self.custom_forward_fn(model)
            else:
                self.custom_forward_fn(model)
    model.load_state_dict(sd)
    context_to_use.disable_trace_dynamic_graph()
    return context_to_use.graph