78ba1c8cfcd2c31f1e1b007d2e20e771049e869b,tensorboardX/graph.py,,graph,#Any#Any#,35

Before Change

    make_list_of_nodes(lastVar.grad_fn)
    for node in list_of_nodes:
        // shape = TensorShapeProto(dim=[TensorShapeProto.Dim(size=i) for i in node["attr.shape"]])  ugly...
        shape_str = str(node["attr.shape"]).encode(encoding="utf_8")
        nodes.append(
            NodeDef(name=node["name"], op=node["op"], input=node["inputs"],
                    attr={"shape": AttrValue(s=shape_str)}))  // , "T":AttrValue(type="DT_FLOAT")}))
    return GraphDef(node=nodes, versions=VersionDef(producer=22))

After Change

def graph(model, args):
    with torch.onnx.set_training(model, False):
        trace, _ = torch.jit.trace(model, args)
    torch.onnx._optimize_trace(trace, False)
    graph = trace.graph()
    print(graph)
    list_of_nodes = parse(graph)

Instances