ONNX export failing with INTERNAL ASSERT failure to do with node->kind()

1

I’m trying to ONNX export a yolo model with dynamic input shape while not having to apply grid offsets and anchors post processing. I’m using torch.jit.script to help ONNX export do this.
The following function works fine during normal execution, but fails during ONNX export with the error:

RuntimeError: !node->kind().is_aten() && !node->kind().is_prim() && !node->kind().is_attr() INTERNAL ASSERT FAILED at "/pytorch/torch/csrc/jit/serialization/export.cpp":453, please report a bug to PyTorch. 

@torch.jit.script
def yolo_predict(pred, anchors, stride: int, scale_x_y: float):
    #transpose
    nA = anchors.size(0)
    nB, nC, nH, nW = pred.size()
    nattr = int(nC / nA)
    pred = pred.view(nB, nA, nattr, nH, nW).permute(0, 1, 3, 4, 2).contiguous()

    #make grid
    grid_x = torch.arange(nW).view(1, 1, 1, nW)
    grid_y = torch.arange(nH).view(1, 1, nH, 1)
    anchor_w = anchors[:, 0].view(1, nA, 1, 1).float() / (nW * stride)
    anchor_h = anchors[:, 1].view(1, nA, 1, 1).float() / (nH * stride)

    #apply grid
    x = pred[..., 0]
    y = pred[..., 1]
    w = pred[..., 2]
    h = pred[..., 3]
    p = pred[..., 4:]
    bx = ((torch.sigmoid(x) - 0.5) * scale_x_y + 0.5 + grid_x) / nW
    by = ((torch.sigmoid(y) - 0.5) * scale_x_y + 0.5 + grid_y) / nH
    bw = torch.exp(w) * anchor_w
    bh = torch.exp(h) * anchor_h
    bp = torch.sigmoid(p)
    preds = torch.stack([bx, by, bw, bh], -1)
    preds = torch.cat([preds, bp], -1)
    preds = preds.view(nB, -1, nattr)

    return preds

Any ideas what’s causing this error?

2

Could you create an issue on GitHub please, so that we can track this error?

3

I have already, but maintainers are being a bit slow at responding. The issue is https://github.com/pytorch/pytorch/issues/45816