MaskedComputationLayer, fix get_rec_initial_extra_outputs call

rwth-i6 · Mar 8, 2022 · 206577b · 206577b
1 parent 0e5f7e4
commit 206577b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/returnn/tf/layers/rec.py b/returnn/tf/layers/rec.py
@@ -7981,9 +7981,11 @@ def get_rec_initial_extra_outputs(cls, batch_dim, rec_layer, **kwargs):
     assert isinstance(output, Data)
     assert issubclass(layer_class, LayerBase)
     with layer_class.cls_setup_scope(**kwargs):
-      d = layer_class.get_rec_initial_extra_outputs(batch_dim=batch_dim, rec_layer=rec_layer, **layer_desc)
+      output = output.copy_as_batch_major()
+      d = layer_class.get_rec_initial_extra_outputs(
+        batch_dim=batch_dim, rec_layer=rec_layer, output=output, **layer_desc)
       initial_out = layer_class.get_rec_initial_output(
-        batch_dim=batch_dim, rec_layer=rec_layer, output=output.copy_as_batch_major(), **layer_desc)
+        batch_dim=batch_dim, rec_layer=rec_layer, output=output, **layer_desc)
       assert "_output" not in d
       d["_output"] = initial_out
       return d