Update the batch_size in the test part

[lldelisle]

Hi, Here is a proposition which helped me to solve the CUDA memory error.

[LukasMahieu]

Thanks!