When I run your code, there is an error 'no module named 'torchvision.layers'. and I can not use the link
pip install git+git://github.com/pytorch/vision.git@24577864e92b72f7066e1ed16e978e873e19d13d.
The solution of
sudo apt-get install git-core
pip install git+git://github.com/pytorch/vision.git@24577864e92b72f7066e1ed16e978e873e19d13d
doesn't work.