fails with Conv3d

Thanks for open sourcing this repo!

I tried it on my code, but it breaks when used on 3D convolutional layer, because they have 5d weights.
I saw that there is a merge_dims arguments and it does solve the problem, but I still wanted to ask if I should use this or what implications it might have to set merge_dim=True