rocm build (#15981)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/15981

caffe2/operators/unique_ops.cu translated to caffe2/operators/hip/unique_ops.hip breaks rocm build

Reviewed By: BIT-silence

Differential Revision: D13646129

fbshipit-source-id: 900a14e14216686ec4560b30df2eabbd7ec2ff91

diff --git a/caffe2/operators/unique_ops.cu b/caffe2/operators/unique_ops.cu
index e8d2e5d..88c6aad 100644 (file)
--- a/caffe2/operators/unique_ops.cu
+++ b/caffe2/operators/unique_ops.cu
@@ -69,7 +69,7 @@ bool UniqueOp<CUDAContext>::DoRunWithType() {
   }
 
   const T* input = inputTensor.template data<T>();
-  ReinitializeTensor(&thrust_unique_buffer_, {N}, at::dtype<T>().device(Context::GetDeviceType()));
+  ReinitializeTensor(&thrust_unique_buffer_, {N}, at::dtype<T>().device(CUDA));
   auto* buffer = thrust_unique_buffer_.template mutable_data<T>();
   context_.CopyItemsSameDevice(inputTensor.meta(), N, input, buffer);