Fix errors in TF YOLO example (NVIDIA#3379)

Update from `experimental_prefetch_to_device` to `experimental_fetch_to_device` for tf.distribute InputOptions that was changed in more recent TF versions. Remove fixed version from tensorflow-addons requirement - the package is already installed in the DLFW container, so it should be now no-op instead of installing old version from pip. Signed-off-by: Krzysztof Lecki <[email protected]>
AbhijitManepatil · Sep 29, 2021 · 68d73a3 · 68d73a3
1 parent 9e45c07
commit 68d73a3
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/docs/examples/use_cases/tensorflow/yolov4/requirements.txt b/docs/examples/use_cases/tensorflow/yolov4/requirements.txt
@@ -1,4 +1,4 @@
 Cython
 matplotlib
-tensorflow-addons==0.11.2
+tensorflow-addons
 pycocotools
diff --git a/docs/examples/use_cases/tensorflow/yolov4/src/train.py b/docs/examples/use_cases/tensorflow/yolov4/src/train.py
@@ -129,7 +129,7 @@ def dataset_fn(input_context):
 
     input_options = tf.distribute.InputOptions(
         experimental_place_dataset_on_device = True,
-        experimental_prefetch_to_device = False,
+        experimental_fetch_to_device = False,
         experimental_replication_mode = tf.distribute.InputReplicationMode.PER_REPLICA)
 
     dataset = strategy.distribute_datasets_from_function(