bmm vs cuda implementation

[kxhit]

Hi, thanks for opening the source code!

I'm just curious about the acceleration speed of the cuda implementation compared to pytorch.bmm operation, if the input to each MLPs are equal. The test() part code in multi_module.py cannot run successfully due to some flags, and I have no idea how to measure the speed of the cuda implementation against bmm. Could you please give me some guidance? Thanks!