ModuleNotFoundError: No module named 'torch.fx'

[shmorish]

hi

I ran the code "train.py" on my server and got this error.
I searched and found that it was added in PyTorch 1.8.0, so I upgraded PyTorch.
However, I got another error that 「ModuleNotFoundError: No module named 'torch.ao'」

Would you please tell me how to fix the problem?

(DRNet2) shoto@istlab-System-Product-Name:~/DRNet$ sudo python train.py
['/mnt/shoto/DRNet', '/mnt/shoto/anaconda3/envs/DRNet2/lib/python37.zip', '/mnt/shoto/anaconda3/envs/DRNet2/lib/python3.7', '/mnt/shoto/anaconda3/envs/DRNet2/lib/python3.7/lib-dynload', '/mnt/shoto/.local/lib/python3.7/site-packages', '/mnt/shoto/anaconda3/envs/DRNet2/lib/python3.7/site-packages']
/mnt/shoto/.local/lib/python3.7/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: libtorch_cuda_cu.so: cannot open shared object file: No such file or directory
  warn(f"Failed to load image Python extension: {e}")
Traceback (most recent call last):
  File "train.py", line 6, in <module>
    import datasets
  File "/mnt/shoto/DRNet/datasets/__init__.py", line 7, in <module>
    import misc.transforms as own_transforms
  File "/mnt/shoto/DRNet/misc/transforms.py", line 10, in <module>
    from torchvision.transforms import functional as TrF
  File "/mnt/shoto/.local/lib/python3.7/site-packages/torchvision/__init__.py", line 7, in <module>
    from torchvision import models
  File "/mnt/shoto/.local/lib/python3.7/site-packages/torchvision/models/__init__.py", line 2, in <module>
    from .convnext import *
  File "/mnt/shoto/.local/lib/python3.7/site-packages/torchvision/models/convnext.py", line 9, in <module>
    from ..ops.misc import ConvNormActivation
  File "/mnt/shoto/.local/lib/python3.7/site-packages/torchvision/ops/__init__.py", line 18, in <module>
    from .poolers import MultiScaleRoIAlign
  File "/mnt/shoto/.local/lib/python3.7/site-packages/torchvision/ops/poolers.py", line 7, in <module>
    import torch.fx
ModuleNotFoundError: No module named 'torch.fx'