Skip to content

GitLab

Commit d6117ed1 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Fix minor bug in NCF data creation cleanup code. 

PiperOrigin-RevId: 290967096
parent 03038d9c
Show whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
official/recommendation/data_pipeline.py
View file @ d6117ed1
...@@ -425,7 +425,7 @@ class BaseDataConstructor(threading.Thread): ...@@ -425,7 +425,7 @@ class BaseDataConstructor(threading.Thread):
self._shuffle_with_forkpool = not stream_files self._shuffle_with_forkpool = not stream_files
if stream_files: if stream_files:
self._shard_root = epoch_dir or tempfile.mkdtemp(prefix="ncf_") 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: else:
self._shard_root = None self._shard_root = None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment