Torch.jit.script is failing due to unfold?

We use libtorch to deploy, using the following program:

But when it comes to the following unfold :


F.unfold works for me in TorchScript:

def fun(x):
    out = F.unfold(x, kernel_size=2, stride=2, padding=0)
    return out

x = torch.randn(10, 10, 10, 10)
out = fun(x)

fun_scripted = torch.jit.script(fun)
out_scripted = fun_scripted(x)

print((out - out_scripted).abs().max())
# tensor(0.)