def generate_env(
    model: _torch.nn.Module,
    traced_model: _torch.fx.GraphModule,
    input_shape: _Tuple[int],
) -> _Dict[str, _Any]:
    """
    Computes environment dictionary by going through the provided model and
    traced_model with specified input_shape.
    The environment dictionary maps from node name to the actual
    function/method/module/attr/placeholder executed by the node.
    Returns environment dictionary.
    """

    env = dict()
    sample_input = _torch.rand(input_shape)
    count_params = count_model_params(model)
    if count_params > 0 and next(model.parameters()).is_cuda:
        sample_input = sample_input.cuda()
    modules = dict(model.named_modules())

    for node in traced_model.graph.nodes:
        if node.op == "placeholder":
            # Only one graph-level input
            result = sample_input
        elif node.op == "get_attr":
            result = _fetch_attr(model, node.target)
        elif node.op == "call_function":
            result = node.target(
                *_load_arg(node.args, env), **_load_arg(node.kwargs, env)
            )
        elif node.op == "call_method":
            self_obj, *args = _load_arg(node.args, env)
            kwargs = _load_arg(node.kwargs, env)
            result = getattr(self_obj, node.target)(*args, **kwargs)
        elif node.op == "call_module":
            ld_args = _load_arg(node.args, env)
            result = modules[node.target](
                *_load_arg(node.args, env), **_load_arg(node.kwargs, env)
            )
        env[node.name] = result
    return env

    Computes environment dictionary by going through the provided model and
    traced_model with specified input_shape.
    The environment dictionary maps from node name to the actual
    function/method/module/attr/placeholder executed by the node.
    Returns environment dictionary.
    
coremltools/optimize/torch/_utils/graph_utils.py