udpate

Wish · Wish · commit b2a832b991a8 · 2021-09-09T13:29:53.000+08:00
diff --git a/src/tensorRT/onnxplugin/plugins/DCNv2.cu b/src/tensorRT/onnxplugin/plugins/DCNv2.cu
@@ -337,13 +337,11 @@ public:
 	cublasHandle_t cublasHandle_ = nullptr;
 	SetupPlugin(DCNv2);
 
-	int initialize() noexcept{
-		cublasCheck(cublasCreate(&cublasHandle_));
-		return 0;
+	virtual void attachToContext(cudnnContext* /*cudnn*/, cublasContext* cublas, nvinfer1::IGpuAllocator* /*allocator*/) noexcept override{
+		cublasHandle_ = cublas;
 	}
 
-	void terminate() noexcept{
-		cublasCheck(cublasDestroy(cublasHandle_));
+	virtual void detachFromContext() noexcept override{
 		cublasHandle_ = nullptr;
 	}
 

Original file line number	Diff line number	Diff line change
`@@ -337,13 +337,11 @@ public:`
`337`	`337`	`cublasHandle_t cublasHandle_ = nullptr;`
`338`	`338`	`SetupPlugin(DCNv2);`
`339`	`339`
`340`		`- int initialize() noexcept{`
`341`		`- cublasCheck(cublasCreate(&cublasHandle_));`
`342`		`- return 0;`
	`340`	`+ virtual void attachToContext(cudnnContext* /cudnn/, cublasContext* cublas, nvinfer1::IGpuAllocator* /allocator/) noexcept override{`
	`341`	`+ cublasHandle_ = cublas;`
`343`	`342`	`}`
`344`	`343`
`345`		`- void terminate() noexcept{`
`346`		`- cublasCheck(cublasDestroy(cublasHandle_));`
	`344`	`+ virtual void detachFromContext() noexcept override{`
`347`	`345`	`cublasHandle_ = nullptr;`
`348`	`346`	`}`
`349`	`347`