title: Keras Subclassing TypeError: tf__call() got multiple values for argument 'training'
body: <p>I can't seem to find the cause of the issue or a similar issue online. I will update as I figure this out.</p>
  <h2><strong>Background</strong></h2>
  <p>As a project, I'm creating a python script using tensorflow and the keras API to build a custom U-net via model subclassing (following, e.g., <a href="https://keras.io/guides/making_new_layers_and_models_via_subclassing/" rel="nofollow noreferrer">this keras tutorial</a>). I've created my own layer classes, a model class, I've set up the parameters, loaded the data etc. and called a <code>model.fit(...)</code>.</p>
  <h2><strong>The Issue</strong></h2>
  <p>I receive the following error and traceback:</p>
  <pre><code>Epoch 1/50
  Traceback (most recent call last):
  
    File &quot;C:\...\path\to\main\code\my_code.py&quot;, line 418, in &lt;module&gt;
      history = model.fit(x=train_ds, validation_data=val_ds, epochs=epochs, verbose=1)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\keras\engine\training.py&quot;, line 108, in _method_wrapper
      return method(self, *args, **kwargs)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\keras\engine\training.py&quot;, line 1098, in fit
      tmp_logs = train_function(iterator)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\eager\def_function.py&quot;, line 780, in __call__
      result = self._call(*args, **kwds)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\eager\def_function.py&quot;, line 823, in _call
      self._initialize(args, kwds, add_initializers_to=initializers)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\eager\def_function.py&quot;, line 697, in _initialize
      *args, **kwds))
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\eager\function.py&quot;, line 2855, in _get_concrete_function_internal_garbage_collected
      graph_function, _, _ = self._maybe_define_function(args, kwargs)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\eager\function.py&quot;, line 3213, in _maybe_define_function
      graph_function = self._create_graph_function(args, kwargs)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\eager\function.py&quot;, line 3075, in _create_graph_function
      capture_by_value=self._capture_by_value),
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\framework\func_graph.py&quot;, line 986, in func_graph_from_py_func
      func_outputs = python_func(*func_args, **func_kwargs)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\eager\def_function.py&quot;, line 600, in wrapped_fn
      return weak_wrapped_fn().__wrapped__(*args, **kwds)
  
    File &quot;C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\framework\func_graph.py&quot;, line 973, in wrapper
      raise e.ag_error_metadata.to_exception(e)
  
  TypeError: in user code:
  
      C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\keras\engine\training.py:806 train_function  *
          return step_function(self, iterator)
  
      TypeError: tf__call() got multiple values for argument 'training'
  </code></pre>
  <p>The error occurs in the keras training.py script. However I'm new to object-oriented programming, tensorflow, and keras, and so my first thought is that the issue comes from either the way I'm defining the arguments of my class methods or the way I'm using the tensorflow toolset. Below is a simplified version of my code that may help:</p>
  <pre><code>import numpy as np
  import tensorflow as tf
  from tensorflow.keras.layers import Conv2D, Layer, MaxPool2D, UpSampling2D, Input, GaussianNoise 
  from tensorflow.keras.layers import Softmax, LeakyReLU, ReLU, Concatenate, BatchNormalization
  from tensorflow.keras.models import Model
  from tensorflow.keras.optimizers import Adam
  import os
  import pickle
  
  tf.config.experimental.list_physical_devices('GPU')
  
  
  # === Define layer and model classes ===
  
  # ... define custom subclasses that inherit keras.layer.Layer
  # ... example:
  
  class sampleLayer2(Layer):
     def __init__(self, name='layer1', **kwargs):
        super().__init__()
        self.sample_layer1 = sampleLayer1(...)
        self.relu = ReLU()
        self.batch_norm = BatchNormalization()
  
     def call(self, layer_in, training=False):
        x = self.sample_layer1(layer_in)
        x = self.batch_norm(inputs=x)
        x = self.relu(x)
        # ...
        return x
  
  # ...
  # ... other custom subclasses similar to above example...
  # ...
  
  # define model:
  class UNet(Model):
     def __init__(self, input_shape, name='unet', **kwargs):
        super().__init__()
        self.inputs = Input(input_shape)
        self.sample_layera = sampleLayerA()
        self.sample_layerb = sampleLayerB()
        self.sample_layerc = sampleLayerC()
        self.sample_layerd = sampleLayerD()
        self.sample_layere = sampleLayerE()
        # ... other layers that adjust depth ...
        self.gaussian = GaussianNoise(.12)
        self.pool = MaxPool2D(pool_size=(2, 2), strides=(2, 2)) # reduce feature dimensions
        self.upsample = UpSampling2D(size=(2, 2)) # increase feature dimensions
        self.batch_norm = BatchNormalization()
        self.concat = Concatenate()
        self.softmax = Softmax()
        
     def call(self, training=False):
        a = self.inputs
        a = self.gaussian(a) 
        a = self.sample_layera(a) # begin encoder
        b = self.pool(a)
        b = self.sample_layerb(b)
        c = self.pool(b)
        c = self.sample_layerc(c)
        d = self.pool(c)
        d = self.sample_layerd(d)
        e = self.pool(d)
        e = self.sample_layere(e)
        e_up = self.upsample(e)
        # ... layer adjusting depth ...
        d = self.concat([e_up, d])
        d = self.sample_layerd(d)
        d_up = self.upsample(d)
        # ... layer adjusting depth ...
        c = self.concat([d_up, c])
        c = self.sample_layerc(c)
        c_up = self.upsample(c)
        # ... layer adjusting depth ...
        b = self.concat([c_up, b])
        b = self.sample_layerb(b)
        b_up = self.upsample(b)
        # ... layer adjusting depth ...
        a = self.concat([b_up, a])
        a = self.sample_layera(a)
        # ... layer adjusting depth ...
        # ... layer adjusting depth ...
        a = self.batch_norm(inputs=a)
        final_layer = self.softmax(a)
        
        model = Model(inputs=self.inputs, outputs=final_layer)
        
        return model
  
  
  # === Define Parameters ===
  
  np.random.seed(27)
  tf.random.set_seed(27)
  
  img_height = ## some number
  img_width = ##
  input_depth = ##
  epochs = 50
  learning_rate = 5e-4 
  batch_size_train = 32 
  batch_size_val = 128 
  opt = Adam(lr=learning_rate, beta_1=.9, beta_2=.999, epsilon=1e-8)
  
  
  # === Set up Datasets ===
  
  # ...
  # ... define data directories
  # ...
  
  # ...
  # ... define generators that ingest and process input and output data
  # ...
  
  # ...
  # ... Using generators, create tf.data.Dataset objects for training and validation set ...
  # ... result is train_ds_in, train_ds_out, val_ds_in, val_ds_out ...  
  
  train_ds = tf.data.Dataset.zip((train_ds_in, train_ds_out)) # combine input-output pairs in new dataset
  train_ds = train_ds.shuffle(100).batch(batch_size_train) # set batches and shuffling of the dataset
  
  # ...same for validation data ...
  val_ds = tf.data.Dataset.zip((val_ds_in, val_ds_out))
  val_ds = val_ds.shuffle(100).batch(batch_size_val)
  
  
  # === Compile and fit the model ===
  
  input_size = (None, img_height, img_width, input_depth)
  
  model = UNet(input_size)  
  
  model.compile(optimizer=opt, loss='mae', metrics=['accuracy'])
        
  # we already have batches and shuffling
  history = model.fit(x=train_ds, validation_data=val_ds, epochs=epochs, verbose=1)
  </code></pre>
  <h2><strong>My attempts to solve</strong></h2>
  <p>I tried inserting <code>*args</code> or <code>**kwargs</code> as arguments to both model and layer, <code>__init__</code> and <code>call</code> methods, and no change. For all instances where <code>BatchNormalization()</code> is called, I inserted <code>trainable=training</code> and no change. Finally, I removed the <code>training=None</code> argument from the UNet <code>call</code> method and got a slightly different error:</p>
  <pre><code>TypeError: in user code:
  
      C:\...\anaconda3\envs\my_env\lib\site-packages\tensorflow\python\keras\engine\training.py:806 train_function  *
          return step_function(self, iterator)
  
      TypeError: tf__call() takes 1 positional argument but 2 were given
  </code></pre>
  <p>It seems the problem is in the <code>call</code> method in the UNet class. For reference, I use python version 3.7.7, tensorflow version 2.3.1, tf.keras version 2.4.0 ...</p>
  <p><strong>Is there an obvious issue with my approach that causes this problem?</strong> Any useful tips appreciated. Thanks.</p>