FX.graph tracing computation graph in DGL

It seems that the fx.graph cannot trace the computation graph defined in DGL due to the DGL.heterogenous_graph. Is there any workaround for tracing the computation graph from DGL?