[Kernel][Misc] Use TORCH_LIBRARY instead of PYBIND11_MODULE for custom ops (#5047)

2024-06-09 16:23:30 -04:00
parent 5d7e3d0176
commit 5467ac3196
55 changed files with 833 additions and 451 deletions
--- a/csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu
+++ b/csrc/quantization/marlin/sparse/marlin_24_cuda_kernel.cu
@@ -16,7 +16,7 @@
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
-#include <torch/extension.h>
+#include <torch/all.h>

 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>