added demo tfx images

setup.py

@@ -13,7 +13,8 @@
         'pandas>=0.20.3',
         'numpy>=1.13.3',
         'imageio>=2.3.0',
-        'tensorboardx>=1.7'
+        'tensorboard>=1.14.0',
+        'pytorch>=1.1.0'
     ],
     author_email='will@hacstudios.com',
     url='https://github.com/williamFalcon/test_tube',

test_tube/log.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pandas as pd
 from imageio import imwrite
-from tensorboardX import SummaryWriter
+from torch.utils.tensorboard import SummaryWriter
 import atexit
 
 # constants
@@ -95,7 +95,7 @@ def __init__(
 
         # set the tensorboardx log path to the /tf folder in the exp folder
         logdir = self.get_tensorboardx_path(self.name, self.version)
-        super().__init__(logdir=logdir, *args, **kwargs)
+        super().__init__(log_dir=logdir, *args, **kwargs)
 
         # register on exit fx so we always close the writer
         atexit.register(self.on_exit)
@@ -423,9 +423,10 @@ def __hash__(self):
 
     for n_iter in range(20):
         sleep(0.3)
-        e.log({'xsinx': n_iter * np.sin(n_iter)})
+        e.log({'loss/xsinx': n_iter * np.sin(n_iter)})
         if n_iter % 10 == 0:
             print('saved')
             e.save()
 
+    e.close()
     os._exit(1)