ModuleNotFoundError: No module named 'torch.fx.experimental.shape_inference' #135196
Description
🐛 Describe the bug
Reproducing example:
>>> import torch
>>> from torch.fx.experimental.shape_inference.infer_shape import infer_shape
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'torch.fx.experimental.shape_inference'
Versions
Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] onnx==1.14.1
[pip3] onnxruntime==1.16.3
[pip3] optree==0.12.1
[pip3] torch==2.5.0.dev20240902+cpu
[pip3] torch-tb-profiler==0.4.0
[pip3] torchaudio==2.5.0.dev20240902+cpu
[pip3] torchmetrics==1.4.0.post0
[pip3] torchvision==0.20.0.dev20240902+cpu
[conda] Could not collect