It’s nice that your solution works, but I would not suggest to downgrade to version 0.1.x.
You are missing a lot of nice features. See Pytorch Release Notes 3
One thing you should try when porting your code, is to enable warnings highlighting incompatible code.
Quote from the release notes:

Here is a code snippet that you can add to the top of your scripts.
Adding this code will generate warnings highlighting incompatible code.
Fix your code to no longer generate warnings.
insert this to the top of your scripts (usually main.py 3)


import sys, warnings, traceback, torch
def warn_with_traceback(message, category, filename, lineno, file=None, line=None):
sys.stderr.write(warnings.formatwarning(message, category, filename, lineno, line))
traceback.print_stack(sys._getframe(2))
warnings.showwarning = warn_with_traceback; warnings.simplefilter('always', UserWarning);
torch.utils.backcompat.broadcast_warning.enabled = True
torch.utils.backcompat.keepdim_warning.enabled = True

Once all warnings disappear, you can remove the code snippet.