Try2: Colocate slot variables with their base variable in keras OptimizerV2
Not doing so results in errors because ops such as ResourceApplyRMSProp take base as well as slot variables. Such ops cannot be placed when the variable devices are different. PiperOrigin-RevId: 227060221
Loading
Please sign in to comment