# Buggy code
def local_train(self, file_list, batch_size, epoch=1, kwargs=None):
    """
    Local training for local testing. Must in eager mode.
    Argments:
        batch_size: batch size in training
        epoch: the number of epoch in training
        kwargs: contains a dict of parameters used in training
    """
    optimizer = self._opt_fn()
    for _ in range(epoch):
        for f in file_list:
            with recordio.File(f, "r") as rdio_r:
                reader = rdio_r.get_reader(0, rdio_r.count())
                while True:
                    record_buf = list(itertools.islice(reader, 0, batch_size))
                    if not record_buf:
                        break

                    data, labels = self._input_fn(record_buf)

                    with tf.GradientTape() as tape:
                        inputs = []
                        for f_col in self._feature_columns:
                            inputs.append(data[f_col.key])
                        if len(inputs) == 1:
                            inputs = inputs[0]
                        outputs = self._model.call(inputs, training=True)
                        loss = self._loss(outputs, labels)

                        # Add regularization loss if any.
                        # Note: for distributed training, the regularization loss should
                        #       be divided by the number of contributing workers, which
                        #       might be difficult for elasticdl.
                        if self._model.losses:
                            loss += math_ops.add_n(self._model.losses)
                    grads = tape.gradient(loss, self._model.trainable_variables)
                    optimizer.apply_gradients(
                        zip(grads, self._model.trainable_variables)
                    )
                    print("Loss is ", loss.numpy())


# Traceback (AssertionError)
python -m unittest elasticdl/worker/*_test.py
/usr/local/anaconda3/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.
from ._conv import register_converters as _register_converters
2019-05-21 13:57:07.262725: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA
WARNING:tensorflow:From /usr/local/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/resource_variable_ops.py:642: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.
Instructions for updating:
Colocations handled automatically by placer.
must be str, not NoneType
FLoss is  3.090042
Loss is  1.4608976
Loss is  0.913306
Loss is  0.5969497
Loss is  0.66515267
Loss is  0.3935135
Loss is  0.37774342
Loss is  0.289928
.
======================================================================
FAIL: test_distributed_train (elasticdl.worker.worker_test.WorkerTest)
----------------------------------------------------------------------
Traceback (most recent call last):
File "/Users/l.zou/git/elasticdl/elasticdl/worker/worker_test.py", line 96, in test_distributed_train
self.assertTrue(res)
AssertionError: False is not true

----------------------------------------------------------------------
Ran 2 tests in 0.165s

# Fixed code
def local_train(self, file_list, batch_size, epoch=1, kwargs=None):
    """
    Local training for local testing. Must in eager mode.
    Argments:
        batch_size: batch size in training
        epoch: the number of epoch in training
        kwargs: contains a dict of parameters used in training
    """
    optimizer = self._opt_fn()
    for _ in range(epoch):
        for f in file_list:
            with recordio.File(f, "r", decoder=self._codec.decode) as rdio_r:
                reader = rdio_r.get_reader(0, rdio_r.count())
                while True:
                    record_buf = list(itertools.islice(reader, 0, batch_size))
                    if not record_buf:
                        break

                    data, labels = self._input_fn(record_buf)

                    with tf.GradientTape() as tape:
                        inputs = []
                        for f_col in self._feature_columns:
                            inputs.append(data[f_col.key])
                        if len(inputs) == 1:
                            inputs = inputs[0]
                        outputs = self._model.call(inputs, training=True)
                        loss = self._loss(outputs, labels)

                        # Add regularization loss if any.
                        # Note: for distributed training, the regularization loss should
                        #       be divided by the number of contributing workers, which
                        #       might be difficult for elasticdl.
                        if self._model.losses:
                            loss += math_ops.add_n(self._model.losses)
                    grads = tape.gradient(loss, self._model.trainable_variables)
                    optimizer.apply_gradients(
                        zip(grads, self._model.trainable_variables)
                    )
                    print("Loss is ", loss.numpy())