How to iterate over a tensor of unknown size in the call method of a custom layer

Question

The following code contains a runtime error.
Trying to keep things not too complicated here's what my custom layer looks like -

class MyCustomLayer(tf.keras.layers.Layer):
    def __init__(self):
        .
        .
        .

    def build(self, input_shape):
        .
        .
        .

    def call(self, input_tensor):
        .
        .
        .
        output = []
        for elem_index in range(int(tf.size(input_tensor))):
            output.append(
                self.transform_element(elem_index, input_tensor, other_stuff)
            )

        output = tf.reshape(output,input_tensor.shape)
        return output

    def transform_element(self, elem_index, input_tensor, other_stuff) -> int:
        # Returns an integer value
        # This function accesses input_tensor based on elem_index
        .
        .
        .

This code basically attempts to create an input_tensor sized tensor output with each element transformed based on its location.

The problem with this code (as far as I debugged) is that input_tensor.shape = (None, int, int, int)
So tensorflow places a Temporary Tensor when tf.size(input_tensorflow) is called.
So the elem_index is a Placeholder Tensor instead of an integer value.
The transform_element function accesses some elements from input_tensor based on elem_index and computes another integer that it returns. The error occurs when transform_element attempts to calculate this integer. It uses a complicated algorithm (numpy operations and other objects) all assuming that it is transforming an integer, not a tensor (or placeholder).

I was getting the following error initially - NotImplementedError: Cannot convert a symbolic Tensor to a numpy array.

I updated tensorflow based on NotImplementedError: Cannot convert a symbolic Tensor (2nd_target:0) to a numpy array

Now I get ValueError inside transform_element at a place where I'm trying to compare two integers.
The code that computes the new integer is definitely correct (exhaustively tested for integers). I don't know what else to do. Been stuck for more than 3 days on this. Thank you for reading this.

Also is there any better way to do this?