Change input amount of points to the network

[dimentary]

Currently the network is hardcoded to support 2048 points as an input, can I get an advice how to update it to support 4096 points as input?

Right now with 4096 I'm getting RuntimeError: cuDNN error: CUDNN_STATUS_INTERNAL_ERROR because of mismatched shapes most probably. I've noticed that in emd_cuda.cu: const int batch = 2048 is hardcoded, also in some python code.

Can you provide me the list of changes how to solve that if possible? Thank you in advance.