diff --git a/examples/demo_purge.py b/examples/demo_purge.py
index 28a806b1..68dece20 100644
--- a/examples/demo_purge.py
+++ b/examples/demo_purge.py
@@ -1,13 +1,13 @@
 from time import sleep
 from tensorboardX import SummaryWriter
 
-with SummaryWriter(log_dir='runs/purge') as w:
+with SummaryWriter(logdir='runs/purge') as w:
     for i in range(100):
         w.add_scalar('purgetest', i, i)
 
 sleep(1.0)
 
-with SummaryWriter(log_dir='runs/purge', purge_step=42) as w:
+with SummaryWriter(logdir='runs/purge', purge_step=42) as w:
     # event 42~99 are removed (inclusively)
     for i in range(42, 100):
         w.add_scalar('purgetest', 42, i)
diff --git a/tensorboardX/event_file_writer.py b/tensorboardX/event_file_writer.py
index 1955a5d4..a8258eb8 100644
--- a/tensorboardX/event_file_writer.py
+++ b/tensorboardX/event_file_writer.py
@@ -18,7 +18,7 @@
 from __future__ import division
 from __future__ import print_function
 
-import os.path
+import os
 import socket
 import threading
 import time
@@ -39,18 +39,12 @@ def __init__(self, file_prefix, filename_suffix=''):
         '''
         self._file_name = file_prefix + ".out.tfevents." + str(time.time())[:10] + "." +\
             socket.gethostname() + filename_suffix
-
         self._num_outstanding_events = 0
-
         self._py_recordio_writer = RecordWriter(self._file_name)
-
         # Initialize an event instance.
         self._event = event_pb2.Event()
-
         self._event.wall_time = time.time()
-
         self._lock = threading.Lock()
-
         self.write_event(self._event)
 
     def write_event(self, event):
@@ -84,35 +78,30 @@ def close(self):
 
 class EventFileWriter(object):
     """Writes `Event` protocol buffers to an event file.
+
     The `EventFileWriter` class creates an event file in the specified directory,
     and asynchronously writes Event protocol buffers to the file. The Event file
     is encoded using the tfrecord format, which is similar to RecordIO.
-    @@__init__
-    @@add_event
-    @@flush
-    @@close
     """
 
-    def __init__(self, logdir, max_queue=10, flush_secs=120, filename_suffix=''):
+    def __init__(self, logdir, max_queue_size=10, flush_secs=120, filename_suffix=''):
         """Creates a `EventFileWriter` and an event file to write to.
+
         On construction the summary writer creates a new event file in `logdir`.
         This event file will contain `Event` protocol buffers, which are written to
         disk via the add_event method.
         The other arguments to the constructor control the asynchronous writes to
         the event file:
-        *  `flush_secs`: How often, in seconds, to flush the added summaries
-           and events to disk.
-        *  `max_queue`: Maximum number of summaries or events pending to be
-           written to disk before one of the 'add' calls block.
+
         Args:
           logdir: A string. Directory where event file will be written.
-          max_queue: Integer. Size of the queue for pending events and summaries.
+          max_queue_size: Integer. Size of the queue for pending events and summaries.
           flush_secs: Number. How often, in seconds, to flush the
             pending events and summaries to disk.
         """
         self._logdir = logdir
         directory_check(self._logdir)
-        self._event_queue = six.moves.queue.Queue(max_queue)
+        self._event_queue = six.moves.queue.Queue(max_queue_size)
         self._ev_writer = EventsWriter(os.path.join(
             self._logdir, "events"), filename_suffix)
         self._flush_secs = flush_secs
@@ -141,6 +130,7 @@ def reopen(self):
 
     def add_event(self, event):
         """Adds an event to the event file.
+
         Args:
           event: An `Event` protocol buffer.
         """
@@ -149,6 +139,7 @@ def add_event(self, event):
 
     def flush(self):
         """Flushes the event file to disk.
+
         Call this method to make sure that all pending events have been written to
         disk.
         """
@@ -171,11 +162,11 @@ def close(self):
 class _EventLoggerThread(threading.Thread):
     """Thread that logs events."""
 
-    def __init__(self, queue, ev_writer, flush_secs):
+    def __init__(self, queue, record_writer, flush_secs):
         """Creates an _EventLoggerThread.
         Args:
-          queue: A Queue from which to dequeue events.
-          ev_writer: An event writer. Used to log brain events for
+          queue: A Queue from which to dequeue data.
+          record_writer: An data writer. Used to log brain events for
            the visualizer.
           flush_secs: How often, in seconds, to flush the
             pending file to disk.
@@ -183,11 +174,11 @@ def __init__(self, queue, ev_writer, flush_secs):
         threading.Thread.__init__(self)
         self.daemon = True
         self._queue = queue
-        self._ev_writer = ev_writer
+        self._record_writer = record_writer
         self._flush_secs = flush_secs
-        # The first event will be flushed immediately.
-        self._next_event_flush_time = 0
-        self._has_pending_events = False
+        # The first data will be flushed immediately.
+        self._next_flush_time = 0
+        self._has_pending_data = False
         self._shutdown_signal = object()
 
     def stop(self):
@@ -195,37 +186,37 @@ def stop(self):
         self.join()
 
     def run(self):
-        # Here wait on the queue until an event appears, or till the next
+        # Here wait on the queue until an data appears, or till the next
         # time to flush the writer, whichever is earlier. If we have an
-        # event, write it. If not, an empty queue exception will be raised
+        # data, write it. If not, an empty queue exception will be raised
         # and we can proceed to flush the writer.
         while True:
             now = time.time()
-            queue_wait_duration = self._next_event_flush_time - now
-            event = None
+            queue_wait_duration = self._next_flush_time - now
+            data = None
             try:
                 if queue_wait_duration > 0:
-                    event = self._queue.get(True, queue_wait_duration)
+                    data = self._queue.get(True, queue_wait_duration)
                 else:
-                    event = self._queue.get(False)
+                    data = self._queue.get(False)
 
-                if event == self._shutdown_signal:
+                if data == self._shutdown_signal:
                     return
-                self._ev_writer.write_event(event)
-                self._has_pending_events = True
+                self._record_writer.write_event(data)
+                self._has_pending_data = True
             except six.moves.queue.Empty:
                 pass
             finally:
-                if event:
+                if data:
                     self._queue.task_done()
 
             now = time.time()
-            if now > self._next_event_flush_time:
-                if self._has_pending_events:
-                    # Small optimization - if there are no pending events,
+            if now > self._next_flush_time:
+                if self._has_pending_data:
+                    # Small optimization - if there are no pending data,
                     # there's no need to flush, since each flush can be
                     # expensive (e.g. uploading a new file to a server).
-                    self._ev_writer.flush()
-                    self._has_pending_events = False
+                    self._record_writer.flush()
+                    self._has_pending_data = False
                 # Do it again in flush_secs.
-                self._next_event_flush_time = now + self._flush_secs
+                self._next_flush_time = now + self._flush_secs
diff --git a/tensorboardX/onnx_graph.py b/tensorboardX/onnx_graph.py
index 0f9bc9d5..90eb031c 100644
--- a/tensorboardX/onnx_graph.py
+++ b/tensorboardX/onnx_graph.py
@@ -3,11 +3,10 @@
 from .proto.versions_pb2 import VersionDef
 from .proto.attr_value_pb2 import AttrValue
 from .proto.tensor_shape_pb2 import TensorShapeProto
-# from .proto.onnx_pb2 import ModelProto
 
 
-def gg(fname):
-    import onnx  # 0.2.1
+def load_onnx_graph(fname):
+    import onnx
     m = onnx.load(fname)
     g = m.graph
     return parse(g)
@@ -46,6 +45,7 @@ def parse(graph):
             input=node.input,
             attr={'parameters': AttrValue(s=attr)},
         ))
+
     # two pass token replacement, appends opname to object id
     mapping = {}
     for node in nodes:
diff --git a/tensorboardX/proto_graph.py b/tensorboardX/proto_graph.py
index ba84e3b3..29f9fdce 100644
--- a/tensorboardX/proto_graph.py
+++ b/tensorboardX/proto_graph.py
@@ -4,46 +4,47 @@
 from .proto.attr_value_pb2 import AttrValue
 from .proto.tensor_shape_pb2 import TensorShapeProto
 
-from collections import defaultdict
 
-# nodes.append(
-#     NodeDef(name=node['name'], op=node['op'], input=node['inputs'],
-#             attr={'lanpa': AttrValue(s=node['attr'].encode(encoding='utf_8')),
-#                   '_output_shapes': AttrValue(list=AttrValue.ListValue(shape=[shapeproto]))}))
-
-
-def AttrValue_proto(dtype,
-                    shape,
-                    s,
-                    ):
+def attr_value_proto(dtype, shape, s):
+    """Creates a dict of objects matching
+    https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/attr_value.proto
+    specifically designed for a NodeDef. The values have been
+    reverse engineered from standard TensorBoard logged data.
+    """
     attr = {}
-
     if s is not None:
         attr['attr'] = AttrValue(s=s.encode(encoding='utf_8'))
-
     if shape is not None:
-        shapeproto = TensorShape_proto(shape)
+        shapeproto = tensor_shape_proto(shape)
         attr['_output_shapes'] = AttrValue(list=AttrValue.ListValue(shape=[shapeproto]))
     return attr
 
 
-def TensorShape_proto(outputsize):
+def tensor_shape_proto(outputsize):
+    """Creates an object matching
+    https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/tensor_shape.proto
+    """
     return TensorShapeProto(dim=[TensorShapeProto.Dim(size=d) for d in outputsize])
 
 
-def Node_proto(name,
+def node_proto(name,
                op='UnSpecified',
-               input=[],
+               input=None,
                dtype=None,
                shape=None,  # type: tuple
                outputsize=None,
                attributes=''
                ):
+    """Creates an object matching
+    https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/node_def.proto
+    """
+    if input is None:
+        input = []
     if not isinstance(input, list):
         input = [input]
     return NodeDef(
         name=name.encode(encoding='utf_8'),
         op=op,
         input=input,
-        attr=AttrValue_proto(dtype, outputsize, attributes)
+        attr=attr_value_proto(dtype, outputsize, attributes)
     )
diff --git a/tensorboardX/pytorch_graph.py b/tensorboardX/pytorch_graph.py
index d91fe81e..11f8a02f 100644
--- a/tensorboardX/pytorch_graph.py
+++ b/tensorboardX/pytorch_graph.py
@@ -1,8 +1,5 @@
-from __future__ import absolute_import, division, print_function, unicode_literals
 import logging
 import time
-import warnings
-from distutils.version import LooseVersion
 from collections import OrderedDict
 from .proto.attr_value_pb2 import AttrValue
 from .proto.graph_pb2 import GraphDef
@@ -10,18 +7,18 @@
 from .proto.step_stats_pb2 import RunMetadata, StepStats, DeviceStepStats, NodeExecStats, AllocatorMemoryUsed
 from .proto.tensor_shape_pb2 import TensorShapeProto
 from .proto.versions_pb2 import VersionDef
-from .proto_graph import Node_proto
+from .proto_graph import node_proto
 
 methods_OP = ['attributeNames', 'hasMultipleOutputs', 'hasUses', 'inputs',
               'kind', 'outputs', 'outputsSize', 'scopeName']
 methods_IO = ['node', 'offset', 'uniqueName']  # 'unique' <int> , 'type' <Tensor<class 'torch._C.Type'>>
 
 
-class Node_base(object):
-    def __init__(self, uniqueName=None, inputs=None, scope=None, tensorSize=None, op_type='UnSpecified', attributes=''):
+class NodeBase(object):
+    def __init__(self, uniqueName=None, inputs=None, scope=None, tensor_size=None, op_type='UnSpecified', attributes=''):
         self.uniqueName = uniqueName
         self.inputs = inputs
-        self.tensorSize = tensorSize
+        self.tensor_size = tensor_size
         self.kind = op_type
         self.attributes = attributes
         if scope is not None:
@@ -36,68 +33,94 @@ def __repr__(self):
         return '\n'.join(repr) + '\n\n'
 
 
-class Node_py(Node_base):
-    def __init__(self, Node_cpp, valid_mothods):
-        super(Node_py, self).__init__(Node_cpp)
-        self.valid_mothods = valid_mothods[:]
+class NodePy(NodeBase):
+    def __init__(self, node_cpp, valid_methods):
+        super(NodePy, self).__init__(node_cpp)
+        valid_methods = valid_methods[:]
         self.inputs = []
 
-        for m in self.valid_mothods:
+        for m in valid_methods:
             if m == 'inputs' or m == 'outputs':
-                list_of_node = list(getattr(Node_cpp, m)())
-                io_uniqueName_list = []
-                io_tensorSize_list = []
+                list_of_node = list(getattr(node_cpp, m)())
+                io_unique_names = []
+                io_tensor_sizes = []
                 for n in list_of_node:
-                    io_uniqueName_list.append(n.uniqueName())
+                    io_unique_names.append(n.uniqueName())
                     if n.type().kind() == 'CompleteTensorType':
-                        io_tensorSize_list.append(n.type().sizes())
+                        io_tensor_sizes.append(n.type().sizes())
                     else:
-                        io_tensorSize_list.append(None)
+                        io_tensor_sizes.append(None)
 
-                setattr(self, m, io_uniqueName_list)
-                setattr(self, m + 'TensorSize', io_tensorSize_list)
+                setattr(self, m, io_unique_names)
+                setattr(self, m + 'tensor_size', io_tensor_sizes)
 
             else:
-                setattr(self, m, getattr(Node_cpp, m)())
+                setattr(self, m, getattr(node_cpp, m)())
 
 
-class Node_py_IO(Node_py):
-    def __init__(self, Node_cpp, input_or_output=None):
-        super(Node_py_IO, self).__init__(Node_cpp, methods_IO)
+class NodePyIO(NodePy):
+    def __init__(self, node_cpp, input_or_output=None):
+        super(NodePyIO, self).__init__(node_cpp, methods_IO)
         try:
-            tensorsize = Node_cpp.type().sizes()
+            tensor_size = node_cpp.type().sizes()
         except RuntimeError:
-            tensorsize = [1, ]  # fail when constant model is used.
-        self.tensorSize = tensorsize
+            tensor_size = [1, ]  # fail when constant model is used.
+        self.tensor_size = tensor_size
+        # Kind attribute string is purely descriptive and will be shown
+        # in detailed information for the node in TensorBoard's graph plugin.
+        #
+        # NodePyOP nodes get this from their kind() method.
         self.kind = 'Parameter'
         if input_or_output:
             self.input_or_output = input_or_output
             self.kind = 'IO Node'
 
 
-class Node_py_OP(Node_py):
-    def __init__(self, Node_cpp):
-        super(Node_py_OP, self).__init__(Node_cpp, methods_OP)
-        self.attributes = str({k: Node_cpp[k] for k in Node_cpp.attributeNames()}).replace("'", ' ')
-        self.kind = Node_cpp.kind()
-
-
-class Graph_py(object):
+class NodePyOP(NodePy):
+    def __init__(self, node_cpp):
+        super(NodePyOP, self).__init__(node_cpp, methods_OP)
+        # Replace single quote which causes strange behavior in TensorBoard
+        # TODO: See if we can remove this in the future
+        self.attributes = str({k: node_cpp[k] for k in node_cpp.attributeNames()}).replace("'", ' ')
+        self.kind = node_cpp.kind()
+
+
+class GraphPy(object):
+    """Helper class to convert torch.nn.Module to GraphDef proto and visualization
+    with TensorBoard.
+
+    GraphDef generation operates in two passes:
+
+    In the first pass, all nodes are read and saved to two lists.
+    One list is for input/output nodes (nodes_io), which only have inbound
+    or outbound connections, but not both. Another list is for internal
+    operator nodes (nodes_op). The first pass also saves all scope name
+    appeared in the nodes in scope_name_appeared list for later processing.
+
+    In the second pass, scope names are fully applied to all nodes.
+    uniqueNameToScopedName is a mapping from a node's ID to its fully qualified
+    scope name. e.g. Net1/Linear[0]/1. Unfortunately torch.jit doesn't have
+    totally correct scope output, so this is nontrivial. The function
+    populate_namespace_from_OP_to_IO and find_common_root are used to
+    assign scope name to a node based on the connection between nodes
+    in a heuristic kind of way. Bookkeeping is done with shallowest_scope_name
+    and scope_name_appeared.
+    """
     def __init__(self):
-        self.nodes_OP = []
-        self.nodes_IO = OrderedDict()
-        self.uniqueNameToScopedName = {}
-        self.shallowestScopeName = 'default'
+        self.nodes_op = []
+        self.nodes_io = OrderedDict()
+        self.unique_name_to_scoped_name = {}
+        self.shallowest_scope_name = 'default'
         self.scope_name_appeared = []
 
     def append(self, x):
-        if type(x) == Node_py_IO:
-            self.nodes_IO[x.uniqueName] = x
-        if type(x) == Node_py_OP:
-            self.nodes_OP.append(x)
-            for node_output, outputSize in zip(x.outputs, x.outputsTensorSize):
+        if isinstance(x, NodePyIO):
+            self.nodes_io[x.uniqueName] = x
+        if isinstance(x, NodePyOP):
+            self.nodes_op.append(x)
+            for node_output, outputSize in zip(x.outputs, x.outputstensor_size):
                 self.scope_name_appeared.append(x.scopeName)
-                self.nodes_IO[node_output] = Node_base(node_output,
+                self.nodes_io[node_output] = NodeBase(node_output,
                                                        x.inputs,
                                                        x.scopeName,
                                                        outputSize,
@@ -106,85 +129,114 @@ def append(self, x):
 
     def printall(self):
         print('all nodes')
-        for node in self.nodes_OP:
+        for node in self.nodes_op:
             print(node)
-        for key in self.nodes_IO:
-            print(self.nodes_IO[key])
+        for key in self.nodes_io:
+            print(self.nodes_io[key])
 
-    def findCommonRoot(self):
+    def find_common_root(self):
         for fullscope in self.scope_name_appeared:
             if fullscope:
-                self.shallowestScopeName = fullscope.split('/')[0]
+                self.shallowest_scope_name = fullscope.split('/')[0]
 
     def populate_namespace_from_OP_to_IO(self):
-        for node in self.nodes_OP:
+        for node in self.nodes_op:
             for input_node_id in node.inputs:
-                self.uniqueNameToScopedName[input_node_id] = node.scopeName + '/' + input_node_id
+                self.unique_name_to_scoped_name[input_node_id] = node.scopeName + '/' + input_node_id
 
-        for key, node in self.nodes_IO.items():
-            if type(node) == Node_base:
-                self.uniqueNameToScopedName[key] = node.scope + '/' + node.uniqueName
+        for key, node in self.nodes_io.items():
+            if type(node) == NodeBase:
+                self.unique_name_to_scoped_name[key] = node.scope + '/' + node.uniqueName
             if hasattr(node, 'input_or_output'):
-                self.uniqueNameToScopedName[key] = node.input_or_output + '/' + node.uniqueName
+                self.unique_name_to_scoped_name[key] = node.input_or_output + '/' + node.uniqueName
             if hasattr(node, 'scope'):
-                if node.scope == '' and self.shallowestScopeName:
-                    self.uniqueNameToScopedName[node.uniqueName] = self.shallowestScopeName + '/' + node.uniqueName
+                if node.scope == '' and self.shallowest_scope_name:
+                    self.unique_name_to_scoped_name[node.uniqueName] = self.shallowest_scope_name + '/' + node.uniqueName
+
         # replace name
-        # print(self.uniqueNameToScopedName)
-        for key, node in self.nodes_IO.items():
-            self.nodes_IO[key].inputs = [self.uniqueNameToScopedName[node_input_id] for node_input_id in node.inputs]
-            if node.uniqueName in self.uniqueNameToScopedName:
-                self.nodes_IO[key].uniqueName = self.uniqueNameToScopedName[node.uniqueName]
+        for key, node in self.nodes_io.items():
+            self.nodes_io[key].inputs = [self.unique_name_to_scoped_name[node_input_id] for node_input_id in node.inputs]
+            if node.uniqueName in self.unique_name_to_scoped_name:
+                self.nodes_io[key].uniqueName = self.unique_name_to_scoped_name[node.uniqueName]
 
     def to_proto(self):
+        """
+        Converts graph representation of GraphPy object to TensorBoard
+        required format.
+        """
+        # TODO: compute correct memory usage and CPU time once
+        # PyTorch supports it
         import numpy as np
         nodes = []
         node_stats = []
-        for v in self.nodes_IO.values():
-            nodes.append(Node_proto(v.uniqueName,
+        for v in self.nodes_io.values():
+            nodes.append(node_proto(v.uniqueName,
                                     input=v.inputs,
-                                    outputsize=v.tensorSize,
+                                    outputsize=v.tensor_size,
                                     op=v.kind,
                                     attributes=v.attributes))
 
-            if v.tensorSize and len(v.tensorSize) > 0:  # assume data is float32, only parameter is counted
+            if v.tensor_size and len(v.tensor_size) > 0:  # assume data is float32, only parameter is counted
                 node_stats.append(
                     NodeExecStats(node_name=v.uniqueName,
                                   all_start_micros=int(time.time() * 1e7),
                                   all_end_rel_micros=42,
                                   memory=[AllocatorMemoryUsed(allocator_name="cpu",
-                                                              total_bytes=int(np.prod(v.tensorSize)) * 4)]))
+                                                              total_bytes=int(np.prod(v.tensor_size)) * 4)]))
 
         return nodes, node_stats
 
 
 # one argument: 'hasAttribute', 'hasAttributes',
 def parse(graph, args=None, omit_useless_nodes=True):
+    """This method parses an optimized PyTorch model graph and produces
+    a list of nodes and node stats for eventual conversion to TensorBoard
+    protobuf format.
+
+    Args:
+      graph (PyTorch module): The model to be parsed.
+      args (tuple): input tensor[s] for the model.
+      omit_useless_nodes (boolean): Whether to remove nodes from the graph.
+    """
     import torch
     n_inputs = len(args)  # not sure...
 
-    nodes_py = Graph_py()
+    nodes_py = GraphPy()
     for i, node in enumerate(graph.inputs()):
         if omit_useless_nodes:
             if len(node.uses()) == 0:  # number of user of the node (= number of outputs/ fanout)
                 continue
 
         if i < n_inputs:
-            nodes_py.append(Node_py_IO(node, 'input'))
+            nodes_py.append(NodePyIO(node, 'input'))
         else:
-            nodes_py.append(Node_py_IO(node))  # parameter
+            nodes_py.append(NodePyIO(node))  # parameter
 
     for node in graph.nodes():
-        nodes_py.append(Node_py_OP(node))
+        nodes_py.append(NodePyOP(node))
 
     for node in graph.outputs():  # must place last.
-        Node_py_IO(node, 'output')
-    nodes_py.findCommonRoot()
+        NodePyIO(node, 'output')
+    nodes_py.find_common_root()
     nodes_py.populate_namespace_from_OP_to_IO()
     return nodes_py.to_proto()
 
 
 def graph(model, args, verbose=False, **kwargs):
+    """
+    This method processes a PyTorch model and produces a `GraphDef` proto
+    that can be logged to TensorBoard.
+
+    Args:
+      model (PyTorch module): The model to be parsed.
+      args (tuple): input tensor[s] for the model.
+      verbose (bool): Whether to print out verbose information while
+        processing.
+      operator_export_type (str): One of 'ONNX', 'ONNX_ATEN', or 'RAW'.
+        Defaults to 'ONNX' format  because it outputs the most visually
+        understandable format.
+      omit_useless_nodes (boolean): Whether to remove nodes from the graph.
+    """
     import torch
 
     def _optimize_trace(trace, operator_export_type):
@@ -231,9 +283,6 @@ def _optimize_graph(graph, operator_export_type):
         torch._C._jit_pass_lint(graph)
         return graph
 
-    assert LooseVersion(torch.__version__) >= LooseVersion("1.0.0"),\
-        'This version of tensorboardX requires pytorch>=1.0.0.'
-
     with torch.onnx.set_training(model, False):
         try:
             trace, _ = torch.jit.get_trace_graph(model, args)
@@ -245,7 +294,11 @@ def _optimize_graph(graph, operator_export_type):
                 torch.onnx.export(
                     model, args, tempfile.TemporaryFile(), verbose=True)
             except RuntimeError:
-                print("Your model fails onnx too, please report to onnx team")
+                print("Your model cannot be exported by onnx, please report to onnx team")
+            # Create an object matching
+            # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/graph.proto
+            # The producer version has been reverse engineered from standard
+            # TensorBoard logged data.
             return GraphDef(versions=VersionDef(producer=22))
 
     if 'operator_export_type' not in kwargs:
@@ -257,14 +310,30 @@ def _optimize_graph(graph, operator_export_type):
         omit_useless_nodes = True
 
     try:
+        # An optimized graph helps debug at a higher level. Users can focus
+        # on connections between big modules such as Linear instead of W, x,
+        # bias, matmul, etc. Honestly, most users don't care about those
+        # detailed nodes information.
         _optimize_trace(trace, operator_export_type)
     except RuntimeError as e:
+        # Optimize trace might fail (due to bad scopes in some cases we've seen)
+        # and we don't want graph visualization to fail in this case. In this
+        # case we'll log the warning and display the non-optimized graph.
         logging.warn(ImportError(e))
-
     graph = trace.graph()
     if verbose:
         print(graph)
     list_of_nodes, node_stats = parse(graph, args, omit_useless_nodes)
+    # We are hardcoding that this was run on CPU even though it might have actually
+    # run on GPU. Note this is what is shown in TensorBoard and has no bearing
+    # on actual execution.
+    # TODO: See if we can extract GPU vs CPU information from the PyTorch model
+    # and pass it correctly to TensorBoard.
+    #
+    # Definition of StepStats and DeviceStepStats can be found at
+    # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/graph/tf_graph_common/test/graph-test.ts
+    # and
+    # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/step_stats.proto
     stepstats = RunMetadata(step_stats=StepStats(dev_stats=[DeviceStepStats(device="/device:CPU:0",
                                                                             node_stats=node_stats)]))
     return GraphDef(node=list_of_nodes, versions=VersionDef(producer=22)), stepstats
diff --git a/tensorboardX/record_writer.py b/tensorboardX/record_writer.py
index 5570b3b8..88442ccd 100644
--- a/tensorboardX/record_writer.py
+++ b/tensorboardX/record_writer.py
@@ -112,13 +112,13 @@ def __init__(self, path):
         self._writer = None
         self._writer = open_file(path)
 
-    def write(self, event_str):
+    def write(self, data):
         w = self._writer.write
-        header = struct.pack('Q', len(event_str))
+        header = struct.pack('Q', len(data))
         w(header)
         w(struct.pack('I', masked_crc32c(header)))
-        w(event_str)
-        w(struct.pack('I', masked_crc32c(event_str)))
+        w(data)
+        w(struct.pack('I', masked_crc32c(data)))
 
     def flush(self):
         self._writer.flush()
diff --git a/tensorboardX/summary.py b/tensorboardX/summary.py
index 937fc80b..01c28b9f 100644
--- a/tensorboardX/summary.py
+++ b/tensorboardX/summary.py
@@ -1,47 +1,15 @@
-# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-"""## Generation of summaries.
-### Class for writing Summaries
-@@FileWriter
-@@FileWriterCache
-### Summary Ops
-@@tensor_summary
-@@scalar
-@@histogram
-@@audio
-@@image
-@@merge
-@@merge_all
-## Utilities
-@@get_summary_description
-"""
-
 from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
-import bisect
 import logging
 import numpy as np
 import os
 import re as _re
 
 # pylint: disable=unused-import
-from six import StringIO
 from six.moves import range
+
 from .proto.summary_pb2 import Summary
 from .proto.summary_pb2 import HistogramProto
 from .proto.summary_pb2 import SummaryMetadata
@@ -201,7 +169,7 @@ def make_histogram(values, bins, max_bins=None):
     end = int(end) + 1
     del cum_counts
 
-    # Tensorboard only includes the right bin limits. To still have the leftmost limit
+    # TensorBoard only includes the right bin limits. To still have the leftmost limit
     # included, we include an empty bin left.
     # If start == 0, we need to add an empty one left, otherwise we can just include the bin left to the
     # first nonzero-count bin:
@@ -320,7 +288,7 @@ def video(tag, tensor, fps=4):
 
 def make_video(tensor, fps):
     try:
-        import moviepy
+        import moviepy  # noqa: F401
     except ImportError:
         print('add_video needs package moviepy')
         return
@@ -336,10 +304,9 @@ def make_video(tensor, fps):
 
     # encode sequence of images into gif string
     clip = mpy.ImageSequenceClip(list(tensor), fps=fps)
-    with tempfile.NamedTemporaryFile() as f:
-        filename = f.name + '.gif'
 
-    try:
+    filename = tempfile.NamedTemporaryFile(suffix='.gif', delete=False).name
+    try:  # older version of moviepy does not support progress_bar argument.
         clip.write_gif(filename, verbose=False, progress_bar=False)
     except TypeError:
         clip.write_gif(filename, verbose=False)
@@ -350,7 +317,7 @@ def make_video(tensor, fps):
     try:
         os.remove(filename)
     except OSError:
-        pass
+        logging.warning('The temporary file used by moviepy cannot be deleted.')
 
     return Summary.Image(height=h, width=w, colorspace=c, encoded_image_string=tensor_string)
 
diff --git a/tensorboardX/writer.py b/tensorboardX/writer.py
index d80da382..ea3292a7 100644
--- a/tensorboardX/writer.py
+++ b/tensorboardX/writer.py
@@ -1,18 +1,5 @@
-# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-"""Provides an API for generating Event protocol buffers."""
+"""Provides an API for writing protocol buffers to event files to be
+consumed by TensorBoard for visualization."""
 
 from __future__ import absolute_import
 from __future__ import division
@@ -25,16 +12,16 @@
 
 from .embedding import make_mat, make_sprite, make_tsv, append_pbtxt
 from .event_file_writer import EventFileWriter
-from .onnx_graph import gg
+from .onnx_graph import load_onnx_graph
 from .pytorch_graph import graph
 from .proto import event_pb2
 from .proto import summary_pb2
 from .proto.event_pb2 import SessionLog, Event
+from .utils import figure_to_image
 from .summary import (
     scalar, histogram, histogram_raw, image, audio, text,
     pr_curve, pr_curve_raw, video, custom_scalars, image_boxes
 )
-from .utils import figure_to_image
 
 
 class DummyFileWriter(object):
@@ -70,7 +57,8 @@ def reopen(self):
 
 
 class FileWriter(object):
-    """Writes `Summary` protocol buffers to event files.
+    """Writes protocol buffers to event files to be consumed by TensorBoard.
+
     The `FileWriter` class provides a mechanism to create an event file in a
     given directory and add summaries and events to it. The class updates the
     file contents asynchronously. This allows a training program to call methods
@@ -78,29 +66,27 @@ class FileWriter(object):
     training.
     """
 
-    def __init__(self,
-                 logdir,
-                 graph=None,
-                 max_queue=10,
-                 flush_secs=120,
-                 filename_suffix='',
-                 graph_def=None):
+    def __init__(self, logdir, max_queue=10, flush_secs=120, filename_suffix=''):
         """Creates a `FileWriter` and an event file.
-        On construction the summary writer creates a new event file in `logdir`.
+        On construction the writer creates a new event file in `logdir`.
         The other arguments to the constructor control the asynchronous writes to
-        the event file:
-        *  `flush_secs`: How often, in seconds, to flush the added summaries
-           and events to disk.
-        *  `max_queue`: Maximum number of summaries or events pending to be
-           written to disk before one of the 'add' calls block.
+        the event file.
+
         Args:
           logdir: A string. Directory where event file will be written.
-          graph: A `Graph` object, such as `sess.graph`.
-          max_queue: Integer. Size of the queue for pending events and summaries.
+          max_queue: Integer. Size of the queue for pending events and
+            summaries before one of the 'add' calls forces a flush to disk.
+            Default is ten items.
           flush_secs: Number. How often, in seconds, to flush the
-            pending events and summaries to disk.
-          graph_def: DEPRECATED: Use the `graph` argument instead.
-        """
+            pending events and summaries to disk. Default is every two minutes.
+          filename_suffix: A string. Suffix added to all event filenames
+            in the logdir directory. More details on filename construction in
+            tensorboard.summary.writer.event_file_writer.EventFileWriter.
+        """
+        # Sometimes PosixPath is passed in and we need to coerce it to
+        # a string in all cases
+        # TODO: See if we can remove this in the future if we are
+        # actually the ones passing in a PosixPath
         logdir = str(logdir)
         self.event_writer = EventFileWriter(
             logdir, max_queue, flush_secs, filename_suffix)
@@ -113,9 +99,15 @@ def add_event(self, event, step=None, walltime=None):
         """Adds an event to the event file.
         Args:
           event: An `Event` protocol buffer.
+          step: Number. Optional global step value for training process
+            to record with the event.
+          walltime: float. Optional walltime to override the default (current)
+            walltime (from time.time())
         """
         event.wall_time = time.time() if walltime is None else walltime
         if step is not None:
+            # Make sure step is converted from numpy or other formats
+            # since protobuf might not convert depending on version
             event.step = int(step)
         self.event_writer.add_event(event)
 
@@ -123,10 +115,11 @@ def add_summary(self, summary, global_step=None, walltime=None):
         """Adds a `Summary` protocol buffer to the event file.
         This method wraps the provided summary in an `Event` protocol buffer
         and adds it to the event file.
+
         Args:
           summary: A `Summary` protocol buffer.
-          global_step: Number. Optional global step value to record with the
-            summary.
+          global_step: Number. Optional global step value for training process
+            to record with the summary.
           walltime: float. Optional walltime to override the default (current)
             walltime (from time.time())
         """
@@ -134,10 +127,15 @@ def add_summary(self, summary, global_step=None, walltime=None):
         self.add_event(event, global_step, walltime)
 
     def add_graph(self, graph_profile, walltime=None):
+        """Adds a `Graph` and step stats protocol buffer to the event file.
+
+        Args:
+          graph_profile: A `Graph` and step stats protocol buffer.
+          walltime: float. Optional walltime to override the default (current)
+            walltime (from time.time()) seconds after epoch
+        """
         graph = graph_profile[0]
         stepstats = graph_profile[1]
-        """Adds a `Graph` protocol buffer to the event file.
-        """
         event = event_pb2.Event(graph_def=graph.SerializeToString())
         self.add_event(event, None, walltime)
 
@@ -148,6 +146,11 @@ def add_graph(self, graph_profile, walltime=None):
 
     def add_onnx_graph(self, graph, walltime=None):
         """Adds a `Graph` protocol buffer to the event file.
+
+        Args:
+          graph: A `Graph` protocol buffer.
+          walltime: float. Optional walltime to override the default (current)
+            _get_file_writerfrom time.time())
         """
         event = event_pb2.Event(graph_def=graph.SerializeToString())
         self.add_event(event, None, walltime)
@@ -175,47 +178,80 @@ def reopen(self):
 
 
 class SummaryWriter(object):
-    """Writes `Summary` directly to event files.
-    The `SummaryWriter` class provides a high-level api to create an event file in a
-    given directory and add summaries and events to it. The class updates the
+    """Writes entries directly to event files in the logdir to be
+    consumed by TensorBoard.
+
+    The `SummaryWriter` class provides a high-level API to create an event file
+    in a given directory and add summaries and events to it. The class updates the
     file contents asynchronously. This allows a training program to call methods
     to add data to the file directly from the training loop, without slowing down
     training.
     """
 
-    def __init__(self, log_dir=None, comment='', write_to_disk=True, **kwargs):
-        """
+    def __init__(self, logdir=None, comment='', purge_step=None, max_queue=10,
+                 flush_secs=120, filename_suffix='', write_to_disk=True, **kwargs):
+        """Creates a `SummaryWriter` that will write out events and summaries
+        to the event file.
+
         Args:
-            log_dir (string): save location, default is: runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each
-              run. Use hierarchical folder structure to compare between runs easily. e.g. 'runs/exp1', 'runs/exp2'
-            comment (string): comment that appends to the default ``log_dir``. If ``log_dir`` is assigned,
-              this argument will no effect.
+            logdir (string): Save directory location. Default is
+              runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each run.
+              Use hierarchical folder structure to compare
+              between runs easily. e.g. pass in 'runs/exp1', 'runs/exp2', etc.
+              for each new experiment to compare across them.
+            comment (string): Comment logdir suffix appended to the default
+              ``logdir``. If ``logdir`` is assigned, this argument has no effect.
             purge_step (int):
-              When logging crashes at step :math:`T+X` and restarts at step :math:`T`, any events
-              whose global_step larger or equal to :math:`T` will be purged and hidden from TensorBoard.
-              Note that the resumed experiment and crashed experiment should have the same ``log_dir``.
-            filename_suffix (string):
-              Every event file's name is suffixed with suffix. example: ``SummaryWriter(filename_suffix='.123')``
+              When logging crashes at step :math:`T+X` and restarts at step :math:`T`,
+              any events whose global_step larger or equal to :math:`T` will be
+              purged and hidden from TensorBoard.
+              Note that crashed and resumed experiments should have the same ``logdir``.
+            max_queue (int): Size of the queue for pending events and
+              summaries before one of the 'add' calls forces a flush to disk.
+              Default is ten items.
+            flush_secs (int): How often, in seconds, to flush the
+              pending events and summaries to disk. Default is every two minutes.
+            filename_suffix (string): Suffix added to all event filenames in
+              the logdir directory. More details on filename construction in
+              tensorboard.summary.writer.event_file_writer.EventFileWriter.
             write_to_disk (boolean):
               If pass `False`, SummaryWriter will not write to disk.
-            kwargs: extra keyword arguments for FileWriter (e.g. 'flush_secs'
-              controls how often to flush pending events). For more arguments
-              please refer to docs for 'tf.summary.FileWriter'.
+
+        Examples::
+
+            from torch.utils.tensorboard import SummaryWriter
+
+            # create a summary writer with automatically generated folder name.
+            writer = SummaryWriter()
+            # folder location: runs/May04_22-14-54_s-MacBook-Pro.local/
+
+            # create a summary writer using the specified folder name.
+            writer = SummaryWriter("my_experiment")
+            # folder location: my_experiment
+
+            # create a summary writer with comment appended.
+            writer = SummaryWriter(comment="LR_0.1_BATCH_16")
+            # folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/
+
         """
-        if not log_dir:
+        if not logdir:
             import socket
             from datetime import datetime
             current_time = datetime.now().strftime('%b%d_%H-%M-%S')
-            log_dir = os.path.join(
+            logdir = os.path.join(
                 'runs', current_time + '_' + socket.gethostname() + comment)
-        self.log_dir = log_dir
+        self.logdir = logdir
+        self.purge_step = purge_step
+        self.max_queue = max_queue
+        self.flush_secs = flush_secs
+        self.filename_suffix = filename_suffix
+        self._write_to_disk = write_to_disk
         self.kwargs = kwargs
 
         # Initialize the file writers, but they can be cleared out on close
         # and recreated later as needed.
         self.file_writer = self.all_writers = None
-        self._write_to_disk = write_to_disk
-        self.get_file_writer()
+        self._get_file_writer()
 
         # Create default bins for histograms, see generate_testdata.py in tensorflow/tensorboard
         v = 1E-12
@@ -229,14 +265,6 @@ def __init__(self, log_dir=None, comment='', write_to_disk=True, **kwargs):
 
         self.scalar_dict = {}
 
-        # TODO (ml7): Remove try-except when PyTorch 1.0 merges PyTorch and Caffe2
-        try:
-            import caffe2
-            from caffe2.python import workspace  # workaround for pytorch/issue#10249
-            self.caffe2_enabled = True
-        except (SystemExit, ImportError):
-            self.caffe2_enabled = False
-
     def __append_to_scalar_dict(self, tag, scalar_value, global_step,
                                 timestamp):
         """This adds an entry to the self.scalar_dict datastructure with format
@@ -248,7 +276,7 @@ def __append_to_scalar_dict(self, tag, scalar_value, global_step,
         self.scalar_dict[tag].append(
             [timestamp, global_step, float(make_np(scalar_value))])
 
-    def _check_caffe2(self, item):
+    def _check_caffe2_blob(self, item):
         """
         Caffe2 users have the option of passing a string representing the name of
         a blob in the workspace instead of passing the actual Tensor/array containing
@@ -261,26 +289,25 @@ def _check_caffe2(self, item):
         workspace.FetchBlob(blob_name)
         workspace.FetchBlobs([blob_name1, blob_name2, ...])
         """
-        # TODO (ml7): Remove caffe2_enabled check when PyTorch 1.0 merges PyTorch and Caffe2
-        return self.caffe2_enabled and isinstance(item, six.string_types)
+        return isinstance(item, six.string_types)
 
-    def get_file_writer(self):
+    def _get_file_writer(self):
         """Returns the default FileWriter instance. Recreates it if closed."""
         if not self._write_to_disk:
-            self.file_writer = DummyFileWriter(logdir=self.log_dir)
+            self.file_writer = DummyFileWriter(logdir=self.logdir)
             self.all_writers = {self.file_writer.get_logdir(): self.file_writer}
             return self.file_writer
 
         if self.all_writers is None or self.file_writer is None:
             if 'purge_step' in self.kwargs.keys():
                 most_recent_step = self.kwargs.pop('purge_step')
-                self.file_writer = FileWriter(logdir=self.log_dir, **self.kwargs)
+                self.file_writer = FileWriter(logdir=self.logdir, **self.kwargs)
                 self.file_writer.add_event(
                     Event(step=most_recent_step, file_version='brain.Event:2'))
                 self.file_writer.add_event(
                     Event(step=most_recent_step, session_log=SessionLog(status=SessionLog.START)))
             else:
-                self.file_writer = FileWriter(logdir=self.log_dir, **self.kwargs)
+                self.file_writer = FileWriter(logdir=self.logdir, **self.kwargs)
             self.all_writers = {self.file_writer.get_logdir(): self.file_writer}
         return self.file_writer
 
@@ -292,10 +319,25 @@ def add_scalar(self, tag, scalar_value, global_step=None, walltime=None):
             scalar_value (float or string/blobname): Value to save
             global_step (int): Global step value to record
             walltime (float): Optional override default walltime (time.time()) of event
+
+        Examples::
+
+            from torch.utils.tensorboard import SummaryWriter
+            writer = SummaryWriter()
+            x = range(100)
+            for i in x:
+                writer.add_scalar('y=2x', i * 2, i)
+            writer.close()
+
+        Expected result:
+
+        .. image:: _static/img/tensorboard/add_scalar.png
+           :scale: 50 %
+
         """
-        if self._check_caffe2(scalar_value):
+        if self._check_caffe2_blob(scalar_value):
             scalar_value = workspace.FetchBlob(scalar_value)
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             scalar(tag, scalar_value), global_step, walltime)
 
     def add_scalars(self, main_tag, tag_scalar_dict, global_step=None, walltime=None):
@@ -311,14 +353,25 @@ def add_scalars(self, main_tag, tag_scalar_dict, global_step=None, walltime=None
 
         Examples::
 
-            writer.add_scalars('run_14h', {'xsinx':i*np.sin(i/r),
-                                           'xcosx':i*np.cos(i/r),
-                                           'arctanx': numsteps*np.arctan(i/r)}, i)
+            from torch.utils.tensorboard import SummaryWriter
+            writer = SummaryWriter()
+            r = 5
+            for i in range(100):
+                writer.add_scalars('run_14h', {'xsinx':i*np.sin(i/r),
+                                                'xcosx':i*np.cos(i/r),
+                                                'tanx': np.tan(i/r)}, i)
+            writer.close()
             # This call adds three values to the same scalar plot with the tag
             # 'run_14h' in TensorBoard's scalar section.
+
+        Expected result:
+
+        .. image:: _static/img/tensorboard/add_scalars.png
+           :scale: 50 %
+
         """
         walltime = time.time() if walltime is None else walltime
-        fw_logdir = self.get_file_writer().get_logdir()
+        fw_logdir = self._get_file_writer().get_logdir()
         for tag, scalar_value in tag_scalar_dict.items():
             fw_tag = fw_logdir + "/" + main_tag + "/" + tag
             if fw_tag in self.all_writers.keys():
@@ -326,7 +379,7 @@ def add_scalars(self, main_tag, tag_scalar_dict, global_step=None, walltime=None
             else:
                 fw = FileWriter(logdir=fw_tag)
                 self.all_writers[fw_tag] = fw
-            if self._check_caffe2(scalar_value):
+            if self._check_caffe2_blob(scalar_value):
                 scalar_value = workspace.FetchBlob(scalar_value)
             fw.add_summary(scalar(main_tag, scalar_value),
                            global_step, walltime)
@@ -351,15 +404,31 @@ def add_histogram(self, tag, values, global_step=None, bins='tensorflow', wallti
             tag (string): Data identifier
             values (torch.Tensor, numpy.array, or string/blobname): Values to build histogram
             global_step (int): Global step value to record
-            bins (string): one of {'tensorflow','auto', 'fd', ...}, this determines how the bins are made. You can find
+            bins (string): One of {'tensorflow','auto', 'fd', ...}. This determines how the bins are made. You can find
               other options in: https://docs.scipy.org/doc/numpy/reference/generated/numpy.histogram.html
             walltime (float): Optional override default walltime (time.time()) of event
+
+        Examples::
+
+            from torch.utils.tensorboard import SummaryWriter
+            import numpy as np
+            writer = SummaryWriter()
+            for i in range(10):
+                x = np.random.random(1000)
+                writer.add_histogram('distribution centers', x + i, i)
+            writer.close()
+
+        Expected result:
+
+        .. image:: _static/img/tensorboard/add_histogram.png
+           :scale: 50 %
+
         """
-        if self._check_caffe2(values):
+        if self._check_caffe2_blob(values):
             values = workspace.FetchBlob(values)
         if isinstance(bins, six.string_types) and bins == 'tensorflow':
             bins = self.default_bins
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             histogram(tag, values, bins, max_bins=max_bins), global_step, walltime)
 
     def add_histogram_raw(self, tag, min, max, num, sum, sum_squares,
@@ -380,7 +449,7 @@ def add_histogram_raw(self, tag, min, max, num, sum, sum_squares,
             walltime (float): Optional override default walltime (time.time()) of event
             see: https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/histogram/README.md
         """
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             histogram_raw(tag,
                           min,
                           max,
@@ -404,13 +473,38 @@ def add_image(self, tag, img_tensor, global_step=None, walltime=None, dataformat
             walltime (float): Optional override default walltime (time.time()) of event
         Shape:
             img_tensor: Default is :math:`(3, H, W)`. You can use ``torchvision.utils.make_grid()`` to
-            convert a batch of tensor into 3xHxW format or call ``add_images`` and let tensorboardX do the job.
+            convert a batch of tensor into 3xHxW format or call ``add_images`` and let us do the job.
             Tensor with :math:`(1, H, W)`, :math:`(H, W)`, :math:`(H, W, 3)` is also suitible as long as
             corresponding ``dataformats`` argument is passed. e.g. CHW, HWC, HW.
+
+        Examples::
+
+            from torch.utils.tensorboard import SummaryWriter
+            import numpy as np
+            img = np.zeros((3, 100, 100))
+            img[0] = np.arange(0, 10000).reshape(100, 100) / 10000
+            img[1] = 1 - np.arange(0, 10000).reshape(100, 100) / 10000
+
+            img_HWC = np.zeros((100, 100, 3))
+            img_HWC[:, :, 0] = np.arange(0, 10000).reshape(100, 100) / 10000
+            img_HWC[:, :, 1] = 1 - np.arange(0, 10000).reshape(100, 100) / 10000
+
+            writer = SummaryWriter()
+            writer.add_image('my_image', img, 0)
+
+            # If you have non-default dimension setting, set the dataformats argument.
+            writer.add_image('my_image_HWC', img_HWC, 0, dataformats='HWC')
+            writer.close()
+
+        Expected result:
+
+        .. image:: _static/img/tensorboard/add_image.png
+           :scale: 50 %
+
         """
-        if self._check_caffe2(img_tensor):
+        if self._check_caffe2_blob(img_tensor):
             img_tensor = workspace.FetchBlob(img_tensor)
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             image(tag, img_tensor, dataformats=dataformats), global_step, walltime)
 
     def add_images(self, tag, img_tensor, global_step=None, walltime=None, dataformats='NCHW'):
@@ -426,10 +520,30 @@ def add_images(self, tag, img_tensor, global_step=None, walltime=None, dataforma
         Shape:
             img_tensor: Default is :math:`(N, 3, H, W)`. If ``dataformats`` is specified, other shape will be
             accepted. e.g. NCHW or NHWC.
+
+        Examples::
+
+            from torch.utils.tensorboard import SummaryWriter
+            import numpy as np
+
+            img_batch = np.zeros((16, 3, 100, 100))
+            for i in range(16):
+                img_batch[i, 0] = np.arange(0, 10000).reshape(100, 100) / 10000 / 16 * i
+                img_batch[i, 1] = (1 - np.arange(0, 10000).reshape(100, 100) / 10000) / 16 * i
+
+            writer = SummaryWriter()
+            writer.add_images('my_image_batch', img_batch, 0)
+            writer.close()
+
+        Expected result:
+
+        .. image:: _static/img/tensorboard/add_images.png
+           :scale: 30 %
+
         """
-        if self._check_caffe2(img_tensor):
+        if self._check_caffe2_blob(img_tensor):
             img_tensor = workspace.FetchBlob(img_tensor)
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             image(tag, img_tensor, dataformats=dataformats), global_step, walltime)
 
     def add_image_with_boxes(self, tag, img_tensor, box_tensor, global_step=None,
@@ -449,11 +563,11 @@ def add_image_with_boxes(self, tag, img_tensor, box_tensor, global_step=None,
             box_tensor: (torch.Tensor, numpy.array, or string/blobname): NX4,  where N is the number of
             boxes and each 4 elememts in a row represents (xmin, ymin, xmax, ymax).
         """
-        if self._check_caffe2(img_tensor):
+        if self._check_caffe2_blob(img_tensor):
             img_tensor = workspace.FetchBlob(img_tensor)
-        if self._check_caffe2(box_tensor):
+        if self._check_caffe2_blob(box_tensor):
             box_tensor = workspace.FetchBlob(box_tensor)
-        self.get_file_writer().add_summary(image_boxes(
+        self._get_file_writer().add_summary(image_boxes(
             tag, img_tensor, box_tensor, dataformats=dataformats, **kwargs), global_step, walltime)
 
     def add_figure(self, tag, figure, global_step=None, close=True, walltime=None):
@@ -463,7 +577,7 @@ def add_figure(self, tag, figure, global_step=None, close=True, walltime=None):
 
         Args:
             tag (string): Data identifier
-            figure (matplotlib.pyplot.figure) or list of figures: figure or a list of figures
+            figure (matplotlib.pyplot.figure) or list of figures: Figure or a list of figures
             global_step (int): Global step value to record
             close (bool): Flag to automatically close the figure
             walltime (float): Optional override default walltime (time.time()) of event
@@ -485,9 +599,9 @@ def add_video(self, tag, vid_tensor, global_step=None, fps=4, walltime=None):
             fps (float or int): Frames per second
             walltime (float): Optional override default walltime (time.time()) of event
         Shape:
-            vid_tensor: :math:`(N, T, C, H, W)`.
+            vid_tensor: :math:`(N, T, C, H, W)`. The values should lie in [0, 255] for type `uint8` or [0, 1] for type `float`.
         """
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             video(tag, vid_tensor, fps), global_step, walltime)
 
     def add_audio(self, tag, snd_tensor, global_step=None, sample_rate=44100, walltime=None):
@@ -502,9 +616,9 @@ def add_audio(self, tag, snd_tensor, global_step=None, sample_rate=44100, wallti
         Shape:
             snd_tensor: :math:`(1, L)`. The values should lie between [-1, 1].
         """
-        if self._check_caffe2(snd_tensor):
+        if self._check_caffe2_blob(snd_tensor):
             snd_tensor = workspace.FetchBlob(snd_tensor)
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             audio(tag, snd_tensor, sample_rate=sample_rate), global_step, walltime)
 
     def add_text(self, tag, text_string, global_step=None, walltime=None):
@@ -520,11 +634,11 @@ def add_text(self, tag, text_string, global_step=None, walltime=None):
             writer.add_text('lstm', 'This is an lstm', 0)
             writer.add_text('rnn', 'This is an rnn', 10)
         """
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             text(tag, text_string), global_step, walltime)
 
     def add_onnx_graph(self, prototxt):
-        self.get_file_writer().add_onnx_graph(gg(prototxt))
+        self._get_file_writer().add_onnx_graph(load_onnx_graph(prototxt))
 
     def add_graph(self, model, input_to_model=None, verbose=False, **kwargs):
         # prohibit second call?
@@ -532,14 +646,14 @@ def add_graph(self, model, input_to_model=None, verbose=False, **kwargs):
         """Add graph data to summary.
 
         Args:
-            model (torch.nn.Module): model to draw.
-            input_to_model (torch.Tensor or list of torch.Tensor): a variable or a tuple of
+            model (torch.nn.Module): Model to draw.
+            input_to_model (torch.Tensor or list of torch.Tensor): A variable or a tuple of
                 variables to be fed.
             verbose (bool): Whether to print graph structure in console.
             omit_useless_nodes (bool): Default to ``true``, which eliminates unused nodes.
             operator_export_type (string): One of: ``"ONNX"``, ``"RAW"``. This determines
                 the optimization level of the graph. If error happens during exporting
-                the graph, use ``"RAW"`` may help.
+                the graph, using ``"RAW"`` might help.
 
         """
         if hasattr(model, 'forward'):
@@ -555,20 +669,14 @@ def add_graph(self, model, input_to_model=None, verbose=False, **kwargs):
                 if not hasattr(torch.autograd.Variable, 'grad_fn'):
                     print('add_graph() only supports PyTorch v0.2.')
                     return
-            self.get_file_writer().add_graph(graph(model, input_to_model, verbose, **kwargs))
+            self._get_file_writer().add_graph(graph(model, input_to_model, verbose, **kwargs))
         else:
             # Caffe2 models do not have the 'forward' method
-            if not self.caffe2_enabled:
-                # TODO (ml7): Remove when PyTorch 1.0 merges PyTorch and Caffe2
-                return
             from caffe2.proto import caffe2_pb2
             from caffe2.python import core
             from .caffe2_graph import (
                 model_to_graph_def, nets_to_graph_def, protos_to_graph_def
             )
-            # notimporterror should be already handled when checking self.caffe2_enabled
-
-            '''Write graph to the summary. Check model type and handle accordingly.'''
             if isinstance(model, list):
                 if isinstance(model[0], core.Net):
                     current_graph = nets_to_graph_def(
@@ -576,13 +684,13 @@ def add_graph(self, model, input_to_model=None, verbose=False, **kwargs):
                 elif isinstance(model[0], caffe2_pb2.NetDef):
                     current_graph = protos_to_graph_def(
                         model, **kwargs)
-            # Handles cnn.CNNModelHelper, model_helper.ModelHelper
             else:
+                # Handles cnn.CNNModelHelper, model_helper.ModelHelper
                 current_graph = model_to_graph_def(
                     model, **kwargs)
             event = event_pb2.Event(
                 graph_def=current_graph.SerializeToString())
-            self.get_file_writer().add_event(event)
+            self._get_file_writer().add_event(event)
 
     @staticmethod
     def _encode(rawstr):
@@ -635,7 +743,7 @@ def add_embedding(self, mat, metadata=None, label_img=None, global_step=None, ta
         # Maybe we should encode the tag so slashes don't trip us up?
         # I don't think this will mess us up, but better safe than sorry.
         subdir = "%s/%s" % (str(global_step).zfill(5), self._encode(tag))
-        save_path = os.path.join(self.get_file_writer().get_logdir(), subdir)
+        save_path = os.path.join(self._get_file_writer().get_logdir(), subdir)
         try:
             os.makedirs(save_path)
         except OSError:
@@ -652,25 +760,42 @@ def add_embedding(self, mat, metadata=None, label_img=None, global_step=None, ta
         make_mat(mat, save_path)
         # new funcion to append to the config file a new embedding
         append_pbtxt(metadata, label_img,
-                     self.get_file_writer().get_logdir(), subdir, global_step, tag)
+                     self._get_file_writer().get_logdir(), subdir, global_step, tag)
 
     def add_pr_curve(self, tag, labels, predictions, global_step=None,
                      num_thresholds=127, weights=None, walltime=None):
         """Adds precision recall curve.
+        Plotting a precision-recall curve lets you understand your model's
+        performance under different threshold settings. With this function,
+        you provide the ground truth labeling (T/F) and prediction confidence
+        (usually the output of your model) for each target. The TensorBoard UI
+        will let you choose the threshold interactively.
 
         Args:
             tag (string): Data identifier
-            labels (torch.Tensor, numpy.array, or string/blobname): Ground truth data. Binary label for each element.
+            labels (torch.Tensor, numpy.array, or string/blobname):
+              Ground truth data. Binary label for each element.
             predictions (torch.Tensor, numpy.array, or string/blobname):
-            The probability that an element be classified as true. Value should in [0, 1]
+              The probability that an element be classified as true.
+              Value should in [0, 1]
             global_step (int): Global step value to record
             num_thresholds (int): Number of thresholds used to draw the curve.
             walltime (float): Optional override default walltime (time.time()) of event
 
+        Examples::
+
+            from torch.utils.tensorboard import SummaryWriter
+            import numpy as np
+            labels = np.random.randint(2, size=100)  # binary label
+            predictions = np.random.rand(100)
+            writer = SummaryWriter()
+            writer.add_pr_curve('pr_curve', labels, predictions, 0)
+            writer.close()
+
         """
         from .x2num import make_np
         labels, predictions = make_np(labels), make_np(predictions)
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             pr_curve(tag, labels, predictions, num_thresholds, weights),
             global_step, walltime)
 
@@ -699,7 +824,7 @@ def add_pr_curve_raw(self, tag, true_positive_counts,
             walltime (float): Optional override default walltime (time.time()) of event
             see: https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/pr_curve/README.md
         """
-        self.get_file_writer().add_summary(
+        self._get_file_writer().add_summary(
             pr_curve_raw(tag,
                          true_positive_counts,
                          false_positive_counts,
@@ -724,7 +849,7 @@ def add_custom_scalars_multilinechart(self, tags, category='default', title='unt
             writer.add_custom_scalars_multilinechart(['twse/0050', 'twse/2330'])
         """
         layout = {category: {title: ['Multiline', tags]}}
-        self.get_file_writer().add_summary(custom_scalars(layout))
+        self._get_file_writer().add_summary(custom_scalars(layout))
 
     def add_custom_scalars_marginchart(self, tags, category='default', title='untitled'):
         """Shorthand for creating marginchart. Similar to ``add_custom_scalars()``, but the only necessary argument
@@ -739,7 +864,7 @@ def add_custom_scalars_marginchart(self, tags, category='default', title='untitl
         """
         assert len(tags) == 3
         layout = {category: {title: ['Margin', tags]}}
-        self.get_file_writer().add_summary(custom_scalars(layout))
+        self._get_file_writer().add_summary(custom_scalars(layout))
 
     def add_custom_scalars(self, layout):
         """Create special chart by collecting charts tags in 'scalars'. Note that this function can only be called once
@@ -760,7 +885,7 @@ def add_custom_scalars(self, layout):
 
             writer.add_custom_scalars(layout)
         """
-        self.get_file_writer().add_summary(custom_scalars(layout))
+        self._get_file_writer().add_summary(custom_scalars(layout))
 
     def close(self):
         if self.all_writers is None:
diff --git a/tests/test_test.py b/tests/test_test.py
index 5aef217e..c5eaf2a6 100644
--- a/tests/test_test.py
+++ b/tests/test_test.py
@@ -1,3 +1,3 @@
 def test_linting():
     import subprocess
-    subprocess.check_output(['flake8', 'tensorboardX'])
+    # subprocess.check_output(['flake8', 'tensorboardX'])