Fix import error

Author: ethanwharris

CHANGELOG.md

@@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Removed
 ### Fixed
 - Fixed a bug where for_steps would sometimes not work as expected if called in the wrong order
+- Fixed a bug where torchbearer installed via pip would crash on import
 
 ## [0.3.1] - 2019-05-24
 ### Added

setup.py

@@ -15,7 +15,7 @@
 setup(
     name='torchbearer',
     version=version_dict['__version__'],
-    packages=['torchbearer', 'torchbearer.metrics', 'torchbearer.callbacks', 'torchbearer.variational', 'tests', 'tests.metrics', 'tests.callbacks', 'tests.variational'],
+    packages=['torchbearer', 'torchbearer.metrics', 'torchbearer.callbacks', 'torchbearer.callbacks.imaging', 'torchbearer.variational', 'tests', 'tests.metrics', 'tests.callbacks', 'tests.callbacks.imaging', 'tests.variational'],
     url='https://github.com/ecs-vlc/torchbearer',
     download_url='https://github.com/ecs-vlc/torchbearer/archive/' + version_dict['__version__'] + '.tar.gz',
     license='GPL-3.0',