Add missing device arg. in __restore_checkpoint()

[DrShushen]

Fixing the issue in ModelTrainer.train() - missing device argument in the self.__restore_checkpoint() call.

Tests passing:

$ python setup.py test
WARNING: Testing via this command is deprecated and will be removed in a future version. Users looking for a generic test entry point independent of test runner are encouraged to use tox. 
running test
running egg_info
writing ptgnn.egg-info/PKG-INFO
writing dependency_links to ptgnn.egg-info/dependency_links.txt
writing requirements to ptgnn.egg-info/requires.txt
writing top-level names to ptgnn.egg-info/top_level.txt
writing manifest file 'ptgnn.egg-info/SOURCES.txt'
running build_ext
ok
ok
ok

----------------------------------------------------------------------
Ran 3 tests in 107.874s

OK

[mmjb]

Thanks for contributing this back!