You have to set it yourself, so if your code is in model.py, running
$ PYTORCH_JIT=0 python model.py

would set the variable (assuming you’re using bash).
As for your code in script, I think it may have something to do with your input shapes, this snippet below works for me (TorchScript is statically typed so the is_tensor checks would also evaluate to true, so they’re removed). It’s slicing data.shape[-3] implies that the inputs need to be at least 4 dimensions.
@torch.jit.script
def normalize(data: torch.Tensor, mean: torch.Tensor,
              std: torch.Tensor) -> torch.Tensor:
    if len(mean) != data.shape[-3] and mean.shape[:2] != data.shape[:2]:
        raise ValueError('mean lenght and number of channels do not match')

    if len(std) != data.shape[-3] and std.shape[:2] != data.shape[:2]:
        raise ValueError('std lenght and number of channels do not match')

    if std.shape != mean.shape:
        raise ValueError('std and mean must have the same shape')

    mean = mean[..., :, None, None].to(data.device)
    std = std[..., :, None, None].to(data.device)

    out = data.sub(mean).div(std)

    return out

print(normalize(torch.ones(2, 2, 2, 2), torch.ones(2, 2, 2, 2), torch.ones(2, 2, 2, 2)))