Trouble initialize Rainbow Agent

```
RainbowAgent(sess=tensorflow.compat.v1.Session(),
num_actions=2,
stack_size=4,
observation_shape=samples[0].shape,
observation_dtype=tensorflow.float32)
```

I get error

```
ValueError: Dimension 0 in both shapes must be equal, but are 32 and 3232. Shapes are [32] and [3232]. for '{{node concat}} = ConcatV2[N=2, T=DT_INT64, Tidx=DT_INT32](strided_slice_4, strided_slice_3, concat/axis)' with input shapes: [32,1], [3232,1], [] and with computed input tensors: input[2] = <1>.
In call to configurable 'DQNAgent' ()
In call to configurable 'RainbowAgent' ()
```

In this part of the code:

`rainbowagent.py` `def _build_target_distribution(self)`

```
batch_indices = tf.range(tf.cast(batch_size, tf.int64))[:, None]
# size of next_qt_argmax: batch_size x 2
batch_indexed_next_qt_argmax = tf.concat(
[batch_indices, next_qt_argmax], axis=1)
```

What am i doing wrong?