def __init__(self, size: tf.Tensor, input_proto_index: tf.Tensor):
    super().__init__(size)
    self._input_proto_index = input_proto_index
def __init__(self, size: tf.Tensor, input_proto_index: tf.Tensor):
    super().__init__(size)
    self._input_proto_index = input_proto_index