message: Set the correct GPU context in the continuation used in Unique.
diff: @@ -32,7 +32,9 @@ limitations under the License.
  #if GOOGLE_CUDA
  #include "tensorflow/core/util/cuda_solvers.h" // For ScratchSpace
  +#include "tensorflow/stream_executor/cuda/cuda_activation.h"
  #elif TENSORFLOW_USE_ROCM
  +#include "tensorflow/core/platform/rocm.h"
  #include "tensorflow/core/util/rocm_solvers.h"
  #endif
  @@ -330,6 +332,9 @@ class UniqueOpGPU : public AsyncOpKernel {
  const GPUDevice& device = context->eigen_gpu_device();
  int64 uniq_size = (*last_idx_host.data()) + 1;
  + se::cuda::ScopedActivateExecutorContext scoped_activation{
  + context->op_device_context()->stream()->parent()};
  +
  Tensor unique_input_inds;
  TIndex* unique_input_inds_ptr = nullptr;
  AllocateTemp(context, uniq_size, &unique_input_inds,
  