# Buggy code
def gen():
    for i, ex in enumerate(features):
        yield (
            {
                "input_ids": ex.input_ids,
                "attention_mask": ex.attention_mask,
                "token_type_ids": ex.token_type_ids,
                "feature_index": i,
                "qas_id": ex.qas_id,
            },
            {
                "start_positions": ex.start_position,
                "end_positions": ex.end_position,
                "cls_index": ex.cls_index,
                "p_mask": ex.p_mask,
                "is_impossible": ex.is_impossible,
            },
        )


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
Traceback (most recent call last):
File "run_glue.py", line 229, in <module>
main()
File "run_glue.py", line 199, in main
compute_metrics=compute_metrics,
File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/transformers/trainer_tf.py", line 48, in __init__
self._setup_training()
File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/transformers/trainer_tf.py", line 58, in _setup_training
self._prepare_dataset()
File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/transformers/trainer_tf.py", line 95, in _prepare_dataset
self.num_train_examples = self.train_dataset.reduce(tf.constant(0), lambda x, _: x + 1).numpy()
File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 1934, in reduce
output_types=structure.get_flat_tensor_types(state_structure)))
File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/ops/gen_dataset_ops.py", line 4661, in reduce_dataset
_ops.raise_from_not_ok_status(e, name)
File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/framework/ops.py", line 6606, in raise_from_not_ok_status
six.raise_from(core._status_to_exception(e.code, message), None)
File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.InvalidArgumentError: TypeError: `generator` yielded an element that could not be converted to the expected type. The expected type was int32, but the yielded element was None.
Traceback (most recent call last):

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 805, in generator_py_func
ret, dtype=dtype.as_numpy_dtype))

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/ops/script_ops.py", line 196, in _convert
result = np.asarray(value, dtype=dtype, order="C")

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/numpy/core/_asarray.py", line 85, in asarray
return array(a, dtype, copy=False, order=order)

TypeError: int() argument must be a string, a bytes-like object or a number, not 'NoneType'


During handling of the above exception, another exception occurred:


Traceback (most recent call last):

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/ops/script_ops.py", line 236, in __call__
ret = func(*args)

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 810, in generator_py_func
"element was %s." % (dtype.name, ret)), sys.exc_info()[2])

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/six.py", line 702, in reraise
raise value.with_traceback(tb)

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/data/ops/dataset_ops.py", line 805, in generator_py_func
ret, dtype=dtype.as_numpy_dtype))

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/tensorflow_core/python/ops/script_ops.py", line 196, in _convert
result = np.asarray(value, dtype=dtype, order="C")

File "/home/ubuntu/anaconda3/envs/tensorflow2_p36/lib/python3.6/site-packages/numpy/core/_asarray.py", line 85, in asarray
return array(a, dtype, copy=False, order=order)

TypeError: `generator` yielded an element that could not be converted to the expected type. The expected type was int32, but the yielded element was None.


[[{{node PyFunc}}]] [Op:ReduceDataset]

# Fixed code
def gen():
    for i, ex in enumerate(features):
        if ex.token_type_ids is None:
            yield (
                {
                    "input_ids": ex.input_ids,
                    "attention_mask": ex.attention_mask,
                    "feature_index": i,
                    "qas_id": ex.qas_id,
                },
                {
                    "start_positions": ex.start_position,
                    "end_positions": ex.end_position,
                    "cls_index": ex.cls_index,
                    "p_mask": ex.p_mask,
                    "is_impossible": ex.is_impossible,
                },
            )
        else:
            yield (
                {
                    "input_ids": ex.input_ids,
                    "attention_mask": ex.attention_mask,
                    "token_type_ids": ex.token_type_ids,
                    "feature_index": i,
                    "qas_id": ex.qas_id,
                },
                {
                    "start_positions": ex.start_position,
                    "end_positions": ex.end_position,
                    "cls_index": ex.cls_index,
                    "p_mask": ex.p_mask,
                    "is_impossible": ex.is_impossible,
                },
            )