I don’t know which additional weights and biases you are introducing, but I guess your new model architecture might not be compatible with the cuDNN definition, so disabling the parameter flattening and using the native approach might be the proper way to implement your custom module.