title: Flag for training and test for custom layer in Keras
body: <p>I want to create a custom keras layer which does something during training and something else for validation or testing.</p>
  
  <pre class="lang-py prettyprint-override"><code>from tensorflow import keras
  K = keras.backend
  from keras.layers import Layer
  import tensorflow as tf
  
  class MyCustomLayer(Layer):
  
      def __init__(self, ratio=0.5, **kwargs):
          self.ratio = ratio
          super(MyCustomLayer, self).__init__(**kwargs)
  
      @tf.function
      def call(self, x, is_training=None):
  
          is_training = K.learning_phase()
          tf.print("training: ", is_training)
          if is_training is 1 or is_training is True:
  
              xs = x * 4
              return xs
          else:
              xs = x*0
              return xs
  
  model = Sequential()
  model.add(Dense(16, input_dim=input_dim))
  model.add(MyCustomLayer(0.5))
  model.add(ReLU())
  model.add(Dense(32, activation='relu'))
  model.add(Dense(16, activation='relu'))
  model.add(Dense(output_dim, activation='softmax', kernel_regularizer=l2(0.01)))
  model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
  
  
  model.fit(X_train, y_train, validation_split=0.05, epochs=5)
  </code></pre>
  
  <p>In the output I always get:</p>
  
  <pre><code>training:  0
  training:  0
  training:  0
  training:  0
  training:  0
  training:  0
  training:  0
  training:  0
  </code></pre>
  
  <p>Does anyone knows how to fix this?</p>