def model_summary(
    model: _torch.nn.Module,
    traced_model: _torch.fx.GraphModule,
    input_shape: _Tuple[int],
) -> ModelSummary:
    """
    Method that takes in a model, traced model and expected input shape, and
    returns a ModelSummary data structure that summarizes the model.
    Note that this method will add a total node entry at the end of the ModelSummary.
    """

    sample_input = _torch.randn(input_shape)
    count_params = count_model_params(model)
    if count_params > 0 and next(model.parameters()).is_cuda:
        sample_input = sample_input.cuda()
    _torch.fx.passes.shape_prop.ShapeProp(traced_model).propagate(sample_input)
    modules = dict(model.named_modules())

    summary = ModelSummary()
    total_param_count = 0
    total_param_size = 0.0
    for node in traced_model.graph.nodes:
        op_kind = None
        if node.op == "placeholder":
            op_kind = "Input"
        elif node.op == "output":
            op_kind = "Output"
        elif node.op == "get_attr":
            op_kind = "Attr"
        elif node.op == "call_module":
            op_kind = modules[node.target].__class__.__name__
        elif node.op == "call_function":
            op_kind = parse_call_function_target(node.target)
        else:
            assert node.op == "call_method", "unsupported node op"
            op_kind = node.target

        param_count = 0
        param_size = 0
        if node.op == "call_module":
            for param_name, param in modules[node.target].named_parameters():
                param_count += volume(param.shape)
                param_size += volume(param.shape) * num_bytes(param.dtype)
        elif node.op == "call_function":
            weight_node = None
            bias_node = None
            if node.target in _LINEAR_CONV_DECONV_FUNCS:
                weight_node = node.args[1]
                bias_node = _fetch_argument(node.args, node.kwargs, None, 2, "bias")
            elif node.target == _torch.nn.functional.batch_norm:
                weight_node = _fetch_argument(node.args, node.kwargs, None, 3, "weight")
                bias_node = _fetch_argument(node.args, node.kwargs, None, 4, "bias")
            elif node.target == _torch.nn.functional.layer_norm:
                weight_node = _fetch_argument(node.args, node.kwargs, None, 2, "weight")
                bias_node = _fetch_argument(node.args, node.kwargs, None, 3, "bias")

            # If weight_node is None, then weight will be None
            # If bias_node is None, then bias will be None
            (weight, bias) = _fetch_func_params(model, weight_node, bias_node)
            if weight is not None:
                param_count += volume(weight.shape)
                param_size += volume(weight.shape) * num_bytes(weight.dtype)
            if bias is not None:
                param_count += volume(bias.shape)
                param_size += volume(bias.shape) * num_bytes(weight.dtype)

        param_size = param_size / 1e6  # 1 MB is 1000 * 1000 bytes

        node_shape = None
        if "tensor_meta" in node.meta:
            if not isinstance(
                node.meta["tensor_meta"], _torch.fx.passes.shape_prop.TensorMetadata
            ):
                # Multiple outputs. Simply use the first output shape
                node_shape = list(node.meta["tensor_meta"][0].shape)
            else:
                node_shape = list(node.meta["tensor_meta"].shape)

        summary.append(node.name, op_kind, node_shape, param_count, param_size)

        total_param_count += param_count
        total_param_size += param_size

    summary.append("Total", None, None, total_param_count, total_param_size)
    return summary

    Method that takes in a model, traced model and expected input shape, and
    returns a ModelSummary data structure that summarizes the model.
    Note that this method will add a total node entry at the end of the ModelSummary.
    
coremltools/optimize/torch/_utils/graph_utils.py