[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.
Loading
Please sign in to comment