I assume you are running this code in PyTorch 1.5.0 (not 0.1.5)?
This code snippet should work and works for me in this version.
Could you post the error message as well as the output of print(torch.__version__), please?