mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-14 20:57:59 +00:00
d4493b27ee
Summary: Make `TORCH_CUDABLAS_CHECK` and `TORCH_CUSOLVER_CHECK` available in custom extensions by exporting the internal functions called by the both macros. Rel: https://github.com/pytorch/pytorch/issues/67073 cc xwang233 ptrblck Pull Request resolved: https://github.com/pytorch/pytorch/pull/67161 Reviewed By: jbschlosser Differential Revision: D31984694 Pulled By: ngimel fbshipit-source-id: 0035ecd1398078cf7d3abc23aaefda57aaa31106
17 lines
422 B
C++
17 lines
422 B
C++
#include <iostream>
|
|
|
|
#include <torch/extension.h>
|
|
#include <ATen/cuda/CUDAContext.h>
|
|
|
|
#include <cublas_v2.h>
|
|
|
|
torch::Tensor noop_cublas_function(torch::Tensor x) {
|
|
cublasHandle_t handle;
|
|
TORCH_CUDABLAS_CHECK(cublasCreate(&handle));
|
|
TORCH_CUDABLAS_CHECK(cublasDestroy(handle));
|
|
return x;
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("noop_cublas_function", &noop_cublas_function, "a cublas function");
|
|
}
|