Internal change

PiperOrigin-RevId: 291819853

Internal change
PiperOrigin-RevId: 291819853
a7d1b2b3 · Bruce Fontaine · A. Unique TensorFlower · 1af7172d · a7d1b2b3
Commit a7d1b2b3 authored Jan 27, 2020 by Bruce Fontaine Committed by A. Unique TensorFlower Jan 27, 2020
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 1 deletion

official/recommendation/data_pipeline.py official/recommendation/data_pipeline.py +2 -1

No files found.
--- a/official/recommendation/data_pipeline.py
+++ b/official/recommendation/data_pipeline.py
@@ -425,6 +425,7 @@ class BaseDataConstructor(threading.Thread):
    self._shuffle_with_forkpool = not stream_files
    if stream_files:
      self._shard_root = epoch_dir or tempfile.mkdtemp(prefix="ncf_")
+      if not create_data_offline:
        atexit.register(tf.io.gfile.rmtree, self._shard_root)
    else:
      self._shard_root = None