Additional colocation options and bugfixes for TensorArray
* colocate_with is now set properly when a TensorArray is passed through a while_loop * added a new argument, "colocate_with_first_write" (default: True; this is the current behavior). If False, the TensorArray is simply placed on the device from the context it's constructed in, and no colocation constraints are added. PiperOrigin-RevId: 157643133
Loading
Please sign in to comment