yitongx

try PyTorch of a lower version (1.0.0 is fine) or replace the 'train_data' with 'data'

Right I didn't think of checking the code. Thanks a lot. And I will pull a request then. :-)