Recently when I write a custom op on gpu in tf, I notice that the program will crash if I try to access the tensor data. I've tried two ways below:
const Tensor& x = context->input(0);
# 1st way
T first_value = x.flat<T>()(0);
# 2nd way
T first_value = ((T*)x.flat<T>().data());
These two ways both crash when I use gpu device, whereas work well in cpu device. I suppose that maybe the tensor data has not been transfered into gpu device yet? so it will access the undefined memory address and thus raise segment fault?
Hope someone can help solve my problem, many thanks!