TastyPancakes/Flux.jl
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

rnnTrainingReserveSize

Browse Source
This commit is contained in:
Mike J Innes 2018-01-26 15:35:14 +00:00
parent 3fb83d642d
commit af0c5523ff
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
src/cuda/cudnn.jl
View File

@ -69,6 +69,13 @@ function rnnWorkspaceSize(r::RNNDesc)
return Int(size[]) return Int(size[])
end end
function rnnTrainingReserveSize(r::RNNDesc)
size = Csize_t[0]
@check ccall((:cudnnGetRNNTrainingReserveSize,libcudnn), cudnnStatus_t, (Ptr{Void}, Ptr{Void}, Cint, Ptr{Ptr{Void}}, Ptr{Csize_t}),
libcudnn_handle[], r, 1, [TensorDesc(r.T, (1,r.input,1))], size)
return Int(size[])
end
function rnnParamSize(r::RNNDesc) function rnnParamSize(r::RNNDesc)
size = Csize_t[0] size = Csize_t[0]
@check ccall((:cudnnGetRNNParamsSize, libcudnn), cudnnStatus_t, (Ptr{Void},Ptr{Void},Ptr{Void},Ptr{Csize_t},Cint), @check ccall((:cudnnGetRNNParamsSize, libcudnn), cudnnStatus_t, (Ptr{Void},Ptr{Void},Ptr{Void},Ptr{Csize_t},Cint),