Fix minor bug in NCF data creation cleanup code.

PiperOrigin-RevId: 290967096

Fix minor bug in NCF data creation cleanup code.
PiperOrigin-RevId: 290967096
d6117ed1 · A. Unique TensorFlower · 03038d9c · d6117ed1
Commit d6117ed1 authored Jan 22, 2020 by A. Unique TensorFlower
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

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

No files found.
--- a/official/recommendation/data_pipeline.py
+++ b/official/recommendation/data_pipeline.py
@@ -425,7 +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_")
-      atexit.register(tf.io.gfile.rmtree, dirname=self._shard_root)
+      atexit.register(tf.io.gfile.rmtree, self._shard_root)
    else:
      self._shard_root = None