HI, when I run test_env.py with my dataset the result is [2, 128, 256] could you please help me what is the reason?