AttributeError: module 'torch.ao' has no attribute 'nn'

[new-a-class]

Thank you for good work.
(unsup_meta) /mnt/workspace/code/ICLR2023/PsCo> CUDA_DEVICES_VISIBLE="4,5" python train.py --model psco --backbone conv4 --prediction --num-shots 1 --dataset omniglot --datadir /mnt/workspace/code/ICLR2023/PsCo/data --logdir logs/omniglot/psco
/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torchvision/image.so: undefined symbol: _ZN3c104impl8GPUTrace13gpuTraceStateE'If you don't plan on using image functionality from torchvision.io, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have libjpeg or libpng installed before building torchvision from source?
warn(
Traceback (most recent call last):
File "/mnt/workspace/code/ICLR2023/PsCo/train.py", line 12, in
import utils
File "/mnt/workspace/code/ICLR2023/PsCo/utils.py", line 17, in
import models
File "/mnt/workspace/code/ICLR2023/PsCo/models.py", line 4, in
import torchvision
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torchvision/init.py", line 6, in
from torchvision import datasets, io, models, ops, transforms, utils
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torchvision/models/init.py", line 17, in
from . import detection, optical_flow, quantization, segmentation, video
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torchvision/models/quantization/init.py", line 3, in
from .mobilenet import *
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torchvision/models/quantization/mobilenet.py", line 1, in
from .mobilenetv2 import * # noqa: F401, F403
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torchvision/models/quantization/mobilenetv2.py", line 5, in
from torch.ao.quantization import DeQuantStub, QuantStub
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torch/ao/quantization/init.py", line 3, in
from .fake_quantize import * # noqa: F403
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torch/ao/quantization/fake_quantize.py", line 8, in
from torch.ao.quantization.observer import (
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torch/ao/quantization/observer.py", line 15, in
from torch.ao.quantization.utils import (
File "/mnt/workspace/workgroup/yfchen/anaconda3/envs/unsup_meta/lib/python3.9/site-packages/torch/ao/quantization/utils.py", line 655, in
) -> torch.ao.nn.quantizable.LSTM:
AttributeError: module 'torch.ao' has no attribute 'nn'