The last dimensions must be the same ???

[Zhaomang Sun]

def test_dct_3d():
for N1 in [2, 5, 32]:
x = np.random.normal(size=(1, N1, N1, N1))
ref = fftpack.dct(x, axis=3, type=2)
ref = fftpack.dct(ref, axis=2, type=2)
ref = fftpack.dct(ref, axis=1, type=2)
act = dct_3d(torch.tensor(x).float()).numpy()
assert np.abs(ref - act).max() < EPS, (ref, act)

def test_idct_3d():
for N1 in [2, 5, 32]:
x = np.random.normal(size=(1, N1, N1, N1))
X = dct_3d(torch.tensor(x).float())
y = idct_3d(X).numpy()
assert np.abs(x - y).max() < EPS, x

In the above code, the last three dims are the same. If they are different, met with a problem:
RuntimeError: size mismatch

[Ziyang Hu]

Please take a look at https://github.com/zh217/torch-dct/blob/master/torch_dct/test/test_dct.py#L111-L120 . That test uses different dimensions and it passed when I just ran it. If problem persists, please provide a failing test.