I'm working with a tf.Variable for which I need to update the values using the scatter_nd_update function. The problem is that my variable has a number of elements equals to 512x512x2048x7 = 3 758 096 384 which is larger than the range converd by int32 (2147483647). And apparently the scatter_nd_update function works only with tf.int32 number and does not work with tf.int64.
Here is an example of the code that produces an overflow error:
import tensorflow as tf
image = tf.Variable(tf.ones((512,512,2048,7), dtype=tf.float32))
update = tf.constant([0.0])
index = tf.constant([[0,0,0,0]], dtype=tf.int64)
image.scatter_nd_update(index, update)
print(f"image shape {tf.shape(image)}")
The error that is raised is the following: F ./tensorflow/core/util/gpu_launch_config.h:129] Check failed: work_element_count > 0 (-536870912 vs. 0) Aborted (core dumped)
Do you have any idea about how this problem can be resolved?