Model.predict() throws error when model definition is subclassed

General Discussion
,
#1

I defined my model using model subclassing which takes in training data in the form [pair1,pair2,label] for training and [pair1,pair2] for prediction. Training works fine however I get an error when I call model.predict() on the prediction data. All data pipelines are created using tf.data.Dataset.

Here is the error message I am getting ;

        return step_function(self, iterator)
    /tmp/ipykernel_33/3417466394.py:16 call  *
        x1,x2 = inputs
    /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:520 __iter__
        self._disallow_iteration()
    /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:513 _disallow_iteration
        self._disallow_when_autograph_enabled("iterating over tf.Tensor`")
   ` /opt/conda/lib/python3.7/site-packages/tensorflow/python/framework/ops.py:491 _disallow_when_autograph_enabled
        " indicate you are trying to use an unsupported feature.".format(task))`

    OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

Below is my model definition;

    def __init__(self,num_classes):
        super(Ricenet,self).__init__()
        self.num_classes = num_classes
        
        self.base_model = tf.keras.applications.efficientnet.EfficientNetB0(
            weights='imagenet', input_shape=[*IMG_SHAPE,3], include_top=False
        )
        self.base_model.trainable = False
        
        self.gap = tf.keras.layers.GlobalAveragePooling2D()
        self.concat = tf.keras.layers.Concatenate(axis=1)
        self.classifier = tf.keras.layers.Dense(num_classes,activation="sigmoid",name="output")
    
    def call(self,inputs):
        x1,x2 = inputs
        x1 = self.base_model(x1)
        x2 = self.base_model(x2)
        
        x1 = self.gap(x1)
        x2 = self.gap(x2)
        
        x = self.concat([x1,x2])
        return self.classifier(x) 
    
    def build_model(self,inputs):
        x1,x2 = inputs
        input1 = tf.keras.Input(shape=x1)
        input2 = tf.keras.Input(shape=x2)
        return tf.keras.Model(inputs=[input1,input2],outputs=self.call([input1,input2]))
#3

I have been able to solve this issue now and I am able to train and predict with no issues. it turns out my way of creating the data using the tf.data.Dataset was wrong. After making that correction everything works as it should be.