Add algo argument
@@ -318,15 +318,16 @@ def set_dropout_descriptor(desc, handle, dropout):
def create_rnn_descriptor(hidden_size, num_layers, dropout_desc,
- input_mode, direction, mode, data_type):
+ input_mode, direction, mode, data_type, algo=None):
desc = Descriptor(cudnn.createRNNDescriptor(),
py_cudnn.destroyRNNDescriptor)
if _cudnn_version >= 6000:
_handle = get_handle()
- _algo = cudnn.CUDNN_RNN_ALGO_STANDARD
+ if algo is None:
+ algo = cudnn.CUDNN_RNN_ALGO_STANDARD
cudnn.setRNNDescriptor_v6(
_handle, desc.value, hidden_size, num_layers, dropout_desc.value,
- input_mode, direction, mode, _algo, data_type)
+ input_mode, direction, mode, algo, data_type)
else:
cudnn.setRNNDescriptor_v5(
desc.value, hidden_size, num_layers, dropout_desc.value,
