Issue with torchvision version in requirements (torch.ao)

[amirrezafateh]

Firstly, thank you for providing this code. I've found it to be very useful for my project.
I encountered an issue while running the code due to torchvision version mismatches. In the requirements, it specifies:
conda install pytorch=1.8.1 torchvision cudatoolkit=11.1 -c pytorch
However, the code is compatible with torchvision version 0.9.1, and using the latest version causes problems (specifically, a missing module error related to torch.ao).
please update the requirements to explicitly mention torchvision version 0.9.1
Thank you again for your valuable contribution!