f742426e2280a5699171ee759b7809617218d505,basenji/blocks.py,,attention,#,333

Before Change

  // batch norm
  if batch_norm:
    z = tf.keras.layers.BatchNormalization(
      momentum=bn_momentum,
      gamma_initializer="zeros",
      fused=True)(z)

  // residual add
  current = tf.keras.layers.Add()([current,z])

After Change

  // batch norm
  if batch_norm:
    if bn_type == "sync":
      bn_layer = tf.keras.layers.experimental.SyncBatchNormalization
    else:
      bn_layer = tf.keras.layers.BatchNormalization
    z = bn_layer(
      momentum=bn_momentum,
      gamma_initializer="zeros")(z)

Instances