bug for torch/quantization/tensor_quant.py

[666DZY666]

https://github.com/NVIDIA/TensorRT-Model-Optimizer/blob/0.31.0/modelopt/torch/quantization/tensor_quant.py#L301

/root/anaconda3/lib/python3.10/site-packages/torch/jit/_script.py", line 1381, in script
fn = torch._C._jit_script_compile(
RuntimeError:
Expression of type | cannot be used in a type expression:
@torch.jit.script
def _fake_tensor_quant_backward(inputs, amax: torch.Tensor | None, grad_outputs):

"amax: torch.Tensor | None" is error

[realAsma]

Which Pytorch are you using?

[666DZY666]

Which Pytorch are you using?

2.1

[realAsma]

I am not able to reproduce this error with torch 2.1 and python 3.10. I used the official pytorch docker for 2.1.
Could you please try torch>=2.4 ?