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?