message: Adding ROCm support for staging ops
diff: @@ -216,7 +216,7 @@ class StageOp : public OpKernel {
  };
  REGISTER_KERNEL_BUILDER(Name("Stage").Device(DEVICE_CPU), StageOp);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("Stage").Device(DEVICE_GPU), StageOp);
  #endif
  #ifdef TENSORFLOW_USE_SYCL
  @@ -249,7 +249,7 @@ class UnstageOp : public OpKernel {
  };
  REGISTER_KERNEL_BUILDER(Name("Unstage").Device(DEVICE_CPU), UnstageOp);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("Unstage").Device(DEVICE_GPU), UnstageOp);
  #endif
  #ifdef TENSORFLOW_USE_SYCL
  @@ -284,7 +284,7 @@ class StagePeekOp : public OpKernel {
  };
  REGISTER_KERNEL_BUILDER(Name("StagePeek").Device(DEVICE_CPU), StagePeekOp);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(
  Name("StagePeek").HostMemory("index").Device(DEVICE_GPU), StagePeekOp);
  #endif
  @@ -314,7 +314,7 @@ class StageSizeOp : public OpKernel {
  };
  REGISTER_KERNEL_BUILDER(Name("StageSize").Device(DEVICE_CPU), StageSizeOp);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("StageSize").HostMemory("size").Device(DEVICE_GPU),
  StageSizeOp);
  #endif
  @@ -339,7 +339,7 @@ class StageClearOp : public OpKernel {
  };
  REGISTER_KERNEL_BUILDER(Name("StageClear").Device(DEVICE_CPU), StageClearOp);
  -#if GOOGLE_CUDA
  +#if GOOGLE_CUDA || TENSORFLOW_USE_ROCM
  REGISTER_KERNEL_BUILDER(Name("StageClear").Device(DEVICE_GPU), StageClearOp);
  #endif
  #ifdef TENSORFLOW_USE_SYCL
  