import tensorflow as tf

with tf.device('/cpu:0'):
    try:
        grad = tf.constant([0, 0, 0, 0], dtype=tf.float64, shape=[4])
        indices = tf.constant([-1], dtype=tf.int64, shape=[1])
        segment_ids = tf.constant([-1], dtype=tf.int64, shape=[1])
        dense_output_dim0 = tf.constant([1], dtype=tf.int32, shape=[])

        tf.raw_ops.SparseSegmentSumGradV2(
            grad=grad,
            indices=indices,
            segment_ids=segment_ids,
            dense_output_dim0=dense_output_dim0,
            name=None
        )
        print("SparseSegmentSumGradV2 executed successfully on CPU")
    except Exception as e:
        print(f"Exception on CPU: {e}")
        
with tf.device('/gpu:0'):
    try:
        grad = tf.constant([0, 0, 0, 0], dtype=tf.float64, shape=[4])
        indices = tf.constant([-1], dtype=tf.int64, shape=[1])
        segment_ids = tf.constant([-1], dtype=tf.int64, shape=[1])
        dense_output_dim0 = tf.constant([1], dtype=tf.int32, shape=[])

        tf.raw_ops.SparseSegmentSumGradV2(
            grad=grad,
            indices=indices,
            segment_ids=segment_ids,
            dense_output_dim0=dense_output_dim0,
            name=None
        )
        print("SparseSegmentSumGradV2 executed successfully on GPU")
    except Exception as e:
        print(f"Exception on GPU: {e}")

Output:


I0000 00:00:1748278378.253343  624834 gpu_device.cc:2018] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 2033 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 4050 Laptop GPU, pci bus id: 0000:01:00.0, compute capability: 8.9
2025-05-27 00:52:58.280430: I tensorflow/core/framework/local_rendezvous.cc:407] Local rendezvous is aborting with status: INVALID_ARGUMENT: Index -1 out of range [0, 1).
Exception on CPU: {{function_node __wrapped__SparseSegmentSumGradV2_device_/job:localhost/replica:0/task:0/device:CPU:0}} Index -1 out of range [0, 1). [Op:SparseSegmentSumGradV2] name: 
SparseSegmentSumGradV2 executed successfully on GPU