Make reader respect shard_id pipeline argument in tf.data.Dataset wit…

…h multiple GPUs example (#1850) Signed-off-by: Janusz Lisiecki <[email protected]>
NVIDIA · Apr 3, 2020 · 03ecab0 · 03ecab0
1 parent 571ba4d
commit 03ecab0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docs/examples/frameworks/tensorflow/tensorflow-dataset-multigpu.ipynb b/docs/examples/frameworks/tensorflow/tensorflow-dataset-multigpu.ipynb
@@ -67,7 +67,7 @@
     "        super(MnistPipeline, self).__init__(\n",
     "            batch_size, num_threads, device_id, seed)\n",
     "        self.reader = ops.Caffe2Reader(\n",
-    "            path=data_path, random_shuffle=True, shard_id=0, num_shards=num_shards)\n",
+    "            path=data_path, random_shuffle=True, shard_id=shard_id, num_shards=num_shards)\n",
     "        self.decode = ops.ImageDecoder(\n",
     "            device='mixed',\n",
     "            output_type=types.GRAY)\n",