code: def unconvert(logits):
      "Move root arcs to diagonal"
      new_shape = list(logits.shape)
      new_shape[1] -= 1
      new_shape[2] -= 1
      new_logits = torch.zeros(
          new_shape, dtype=logits.dtype, device=logits.device
      )
  
      new_logits.fill_(-1e9)
      new_logits[:, :, :] = logits[:, 1:, 1:]
      N = new_logits.size(1)
      new_logits[:, torch.arange(N), torch.arange(N)] = logits[:, 0, 1:]
      return new_logits
  
file_path: pytorch-struct/torch_struct/deptree.py
imports: import torch
  from .helpers import _Struct, Chart
label_desc: Use image node_id 2 for calling a global function with example usage: unconvert(logits) and returns: new_logits