#tf-data Fix the graph mode tests for index flat map.
#tf-data Fix the graph mode tests for index flat map.

Otherwise, I'm getting this error:
```
Tensor("Const:0", shape=(4, 3), dtype=int64) must be from the same graph
 as Tensor("TensorSliceDataset:0", shape=(), dtype=variant)
```
