update

Author: mikel-zhobro

returnn/tf/layers/basic.py

@@ -1219,14 +1219,15 @@ def __init__(self, start, size, min_size=None, **kwargs):
     super(SliceNdLayer2, self).__init__(**kwargs)
     from returnn.tf.util.basic import slice_nd2, DimensionTag
     assert start.output.have_batch_axis() and self.input_data.have_batch_axis()
+    self.start = start
 
     input_data = self.input_data.copy_as_batch_major()
     start = start.output.copy_as_batch_major()
 
     # make sure axis of start are in input
     is_equal_opts = dict(ignore_feature_dim=True, allow_same_spatial_dim=True, broadcast_matches=True)
     for start_axis in range(start.batch_ndim):
-      assert input_data.get_dim_tag(start_axis).is_equal(start.get_dim_tag(start_axis), **is_equal_opts), "the input should hold all axis in start"
+      assert input_data.get_dim_tag(start_axis).is_equal(start.get_dim_tag(start_axis), **is_equal_opts)
 
     # Handle the case when layer is pulled out of rec loop but the input hasn't change
     if self.optimized_out_of_loop_and_unchanged_input(input_data, start):
@@ -1261,12 +1262,12 @@ def __init__(self, start, size, min_size=None, **kwargs):
     self.output.placeholder = slices
 
   @classmethod
-  def optimized_out_of_loop_and_unchanged_input(cls, input, start):
+  def optimized_out_of_loop_and_unchanged_input(cls, input_data, start):
     """
     :rtype: bool
     The idea is to check that the axis after the last common axis is a feature axis instead of spatial.
     """
-    return input.get_dim_tag(start.batch_ndim) == input.get_dim_tag(input.get_feature_batch_axes()[0])
+    return input_data.get_dim_tag(start.batch_ndim) == input_data.get_dim_tag(input_data.get_feature_batch_axes()[0])
 
   def get_dep_layers(self):
     """