[JIT] Can't use ndim in script

[jamesr66a]

import torch

def foo(x):
    if x.ndim > 3:
        return torch.neg(x)
    else:
        return x


x = torch.ones(3, 4)
x2 = torch.ones(3, 4, 5, 6)

print(foo(x))
print(foo(x2))

scripted = torch.jit.script(foo)

import torch

def foo(x):
    if x.ndim > 3:
        return torch.neg(x)
    else:
        return x


x = torch.ones(3, 4)
x2 = torch.ones(3, 4, 5, 6)

print(foo(x))
print(foo(x2))

scripted = torch.jit.script(foo)

cc @suo

[eellison]

You can call x.dim() in script

[jamesr66a]

@eellison yeah but this should still work

[eellison]

ndim is deprecated & isn't in the docs. And the error message for this will show dim in the list of similar operators

[jamesr66a]

huh? it's on the docs: