NotFoundError: Could not find valid device for node in TensorFlow for custom loss function on large dataset

Question

Following is my code for the custom loss function

def cost(self, Y, logits):
        loss = tf.math.reduce_sum(
            tf.nn.sparse_softmax_cross_entropy_wit