Fix compatibility with future pytorch

[SkafteNicki]

What does this PR do?

Fixes #1009

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure to update the docs?
• Did you write any new necessary tests?

PR review

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #1011 (41cb93e) into master (1789b17) will decrease coverage by 0%.
The diff coverage is 100%.

@@          Coverage Diff           @@
##           master   #1011   +/-   ##
======================================
- Coverage      95%     95%   -0%     
======================================
  Files         179     179           
  Lines        7652    7656    +4     
======================================
- Hits         7263    7254    -9     
- Misses        389     402   +13