I think I figured it out. For those who interested in getting the “layer indexes”, or to be specific, the “output indexes”, check out this file https://github.com/pytorch/pytorch/blob/98e312cf96f6a5e23933cd8794097063ee3cbc8c/torch/utils/tensorboard/_pytorch_graph.py 126