`stack_compress` Compressed Wrong Dimension

Per demo code in the website, I tried to do the stack compress.

import numpy as np
from cpprb import ReplayBuffer

rb = ReplayBuffer(32,{"obs":{"shape": (16,16,4)}, 'rew': {}, 'done': {}},
                  next_of = "obs", stack_compress = "obs")

rb.add(obs=(np.ones((16,16,4))),
       next_obs=(np.ones((16,16,4))+1),
       rew=1,
       done=0)
rb.add(obs=(np.ones((16,16,4))+1),
       next_obs=(np.ones((16,16,4))+2),
       rew=1,
       done=0)
rb.sample(2)['obs']

And I get the following outputs. Note that the obs could either be of all 1s or all 2s, which does not seem to be the case for data sampled.

array([[[[1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         ...,
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.]],

        [[1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         ...,
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.]],

        [[1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         ...,
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.]],

        ...,

        [[1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         ...,
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.]],

        [[1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         ...,
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.]],

        [[1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         ...,
         [1., 2., 2., 2.],
         [1., 2., 2., 2.],
         [1., 2., 2., 2.]]],


       [[[2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         ...,
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.]],

        [[2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         ...,
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.]],

        [[2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         ...,
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.]],

        ...,

        [[2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         ...,
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.]],

        [[2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         ...,
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.]],

        [[2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         ...,
         [2., 2., 2., 2.],
         [2., 2., 2., 2.],
         [2., 2., 2., 2.]]]], dtype=float32)