编写自定义层时如何训练网络或预测网络

Question

我使用tf.keras编写了一个包含tf.nn.dropout的自定义层。当我通过model.compile编译模型时，如何将训练参数传递或不传递给drpout？ 例如：

class MyFirstLayer(tf.keras.layers.Layer):
       def __init__(self, ):
       ...
       def call(self, inputs,training=0,rate=0.2):
       ...
       if training:
          temp2 = tf.nn.dropout(temp2,rate)
       ...
       return output
Mylayer=MyFisrtLayer()
Mylayer.compile(optmiser=...,loss=...,metrics=...)
Mylayer.fit(...)

由于fit函数同时具有训练和验证过程，如何设置训练参数？