Add variable scopes in self attention encoder
@@ -56,6 +56,7 @@ class SelfAttentionEncoder(Encoder):
for l in range(self.num_layers):
with tf.variable_scope("layer_" + str(l)):
+ with tf.variable_scope("multi_head"):
context = transformer.multi_head_attention(
self.num_heads,
inputs,
@@ -70,6 +71,7 @@ class SelfAttentionEncoder(Encoder):
mode,
dropout=self.dropout)
+ with tf.variable_scope("ffn"):
transformed = transformer.feed_forward(context, self.ffn_inner_dim)
transformed = transformer.add_and_norm(
context,
