提交 f986e0dd authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron

Fix the cudnn context problem.

上级 6b7a3d16
...@@ -102,11 +102,16 @@ setup_ext_cuda(); ...@@ -102,11 +102,16 @@ setup_ext_cuda();
#section support_code_struct #section support_code_struct
PyGpuContextObject *ctx;
cudnnHandle_t APPLY_SPECIFIC(_handle); cudnnHandle_t APPLY_SPECIFIC(_handle);
#section init_code_struct #section init_code_struct
{ {
// We need to keep a reference here to have it available in the destructor.
ctx = PARAMS;
Py_INCREF(ctx);
cuda_enter(PARAMS->ctx); cuda_enter(PARAMS->ctx);
cudnnStatus_t err; cudnnStatus_t err;
APPLY_SPECIFIC(_handle) = NULL; APPLY_SPECIFIC(_handle) = NULL;
...@@ -128,4 +133,7 @@ cudnnHandle_t APPLY_SPECIFIC(_handle); ...@@ -128,4 +133,7 @@ cudnnHandle_t APPLY_SPECIFIC(_handle);
#section cleanup_code_struct #section cleanup_code_struct
cuda_enter(ctx->ctx);
cudnnDestroy(APPLY_SPECIFIC(_handle)); cudnnDestroy(APPLY_SPECIFIC(_handle));
cuda_exit(ctx->ctx);
Py_DECREF((PyObject *)ctx);
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论