Accessing data crash in custom op when using gpu

18 views
Skip to first unread message

cookie

unread,
Nov 2, 2021, 11:21:47 PM11/2/21
to Discuss
Dear community:

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:
```python
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())[0];
```
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!
Reply all
Reply to author
Forward
0 new messages