The current version reads the entire train file into memory and then
loads it with Dataset.from_tensor_slices. This can result in very large
allocations which can either OOM or cause a bad_alloc. To avoid this,
the new version now writes a TFRecord file to the `output_dir`.