update comments

Author: mikel-zhobro

returnn/tf/layers/basic.py

@@ -908,8 +908,8 @@ def __init__(self, start, size, min_size=None, **kwargs):
 
     if seq_lens is not None:
       mask = tf.greater_equal(
-        tf.range(size)[None, :] + tf.expand_dims(start, axis=-1), seq_lens[:, None])  # (B,Tn)
-      mask = expand_multiple_dims(mask, list(range(slice_axis + 2, x.batch_ndim)))  # (B,1,1,..,Tn,1)
+        tf.range(size)[None, :] + tf.expand_dims(start, axis=-1), seq_lens[:, None])  # (B,T1,..,Tn)
+      mask = expand_multiple_dims(mask, list(range(slice_axis + 2, x.batch_ndim)))  # (B,T1,..,Tn,1,..)
       slices = where_bc(mask, tf.zeros_like(slices), slices)
 
     self.output.size_placeholder = x.size_placeholder.copy()