Commit 6e9d3ad2 authored by Derek Murray's avatar Derek Murray Committed by Derek Murray
Browse files

[tf.data] Fix a device placement issue in `prefetch_to_device()`. (#18607)

* [tf.data] Fix a device placement issue in `prefetch_to_device()`.

Previously, the `iterator_get_device()` op was being infeasibly colocated with
both the iterator and placed on the prefetch target device. Move the
construction of that op outside the `with device():` block to fix this.

Also enable the relevant test to run as a CUDA test.

* Import the cuda_py_test rule.
parent 6192a8e5
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment