How to get output shape of each node from a loaded jit traced model?

If a model is traced and sve by torch.jit.trace.Then I load the model just before,the output shapes of nodes in graph are lost, how to get these output shapes of nodes?
Here is an example code:

import torch
import torchvision
from torch._C import _propagate_and_assign_input_shapes
def _model_to_graph(model, args):
    if isinstance(args, torch.Tensor):
        args = (args, )
    graph = model.forward.graph
    method_graph, params = torch._C._jit_pass_lower_graph(graph, model._c)
    in_vars, in_desc = torch.jit._flatten(tuple(args) + tuple(params))
    graph = _propagate_and_assign_input_shapes(  method_graph, tuple(in_vars), False, False)
    return graph
traced_model_savepath = 'traced.pt'
model = torchvision.models.mobilenet.mobilenet_v2(pretrained=True)
dummy_input = torch.rand((1,3,224,224))
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save(traced_model_savepath)
# load the saved traced model from disk
loaded_traced_model = torch.jit.load(traced_model_savepath)
graph= _model_to_graph(loaded_traced_model,  dummy_input)
print(graph)

the output shape of nodes are lost, for example:

  %input.6 : Tensor = aten::_convolution(%626, %251, %276, %627, %628, %629, %277, %630, %274, %277, %277, %278) # /data0/shareVR/pytorch/learn/torch/nn/modules/conv.py:416:0
  %input.104 : Tensor = aten::batch_norm(%input.6, %252, %253, %254, %255, %278, %280, %279, %278) # /data0/shareVR/pytorch/learn/torch/nn/functional.py:2051:0