coreylowman · coreylowman · Jan 14, 2023 · Jan 14, 2023 · Jan 14, 2023
diff --git a/Cargo.toml b/Cargo.toml
@@ -31,7 +31,7 @@ matrixmultiply = { version = "0.3.2", default-features = false }
 zip = { version = "0.6.2", default-features = false, optional = true }
 cblas-sys = { version = "0.1.4", default-features = false, optional = true }
 libc = { version = "0.2", default-features = false, optional = true }
-cudarc = { version = "0.5.1", default-features = false, optional = true }
+cudarc = { version = "0.6.0", default-features = false, optional = true }
 
 [features]
 default = ["std", "numpy"]

diff --git a/src/tensor/cuda/device.rs b/src/tensor/cuda/device.rs
@@ -3,10 +3,8 @@ use crate::tensor::cpu::{Cpu, CpuError};
 use crate::tensor::storage_traits::{DeviceStorage, HasErr};
 
 use cudarc::{
-    blas::CudaBlas,
-    cublas::result::CublasError,
-    device::{BuildError, CudaDevice, CudaDeviceBuilder, CudaSlice},
-    driver::result::DriverError,
+    cublas::{result::CublasError, CudaBlas},
+    driver::{result::DriverError, BuildError, CudaDevice, CudaDeviceBuilder, CudaSlice},
 };
 use std::sync::Arc;
 

diff --git a/src/tensor_ops/abs/cuda_kernel.rs b/src/tensor_ops/abs/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::AbsKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::AbsKernelOp {}
 
 impl UnaryOpCudaKernel for super::AbsKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/abs.ptx"));

diff --git a/src/tensor_ops/add/cuda_kernel.rs b/src/tensor_ops/add/cuda_kernel.rs
@@ -1,7 +1,7 @@
 use crate::tensor_ops::cuda_kernels::{BinaryOpCudaKernel, UnaryOpCudaKernel};
 
-unsafe impl cudarc::device::AsKernelParam for super::ScalarAddKernelOp<f32> {}
-unsafe impl cudarc::device::AsKernelParam for super::BinaryAddKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::ScalarAddKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::BinaryAddKernelOp {}
 
 impl UnaryOpCudaKernel for super::ScalarAddKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/scalar_add.ptx"));

diff --git a/src/tensor_ops/bce/cuda_kernel.rs b/src/tensor_ops/bce/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::BinaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::BCEKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::BCEKernelOp {}
 
 impl BinaryOpCudaKernel for super::BCEKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/bce.ptx"));

diff --git a/src/tensor_ops/broadcast_to/cuda_kernel.rs b/src/tensor_ops/broadcast_to/cuda_kernel.rs
@@ -1,7 +1,7 @@
 use crate::shapes::*;
 use crate::tensor::cuda::{Cuda, CudaArray};
 
-use cudarc::device::{LaunchAsync, LaunchConfig};
+use cudarc::driver::{LaunchAsync, LaunchConfig};
 use std::sync::Arc;
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/broadcast_to.ptx"));

diff --git a/src/tensor_ops/clamp/cuda_kernel.rs b/src/tensor_ops/clamp/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::ClampKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::ClampKernelOp<f32> {}
 
 impl UnaryOpCudaKernel for super::ClampKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/clamp.ptx"));

diff --git a/src/tensor_ops/cos/cuda_kernel.rs b/src/tensor_ops/cos/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::CosKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::CosKernelOp {}
 
 impl UnaryOpCudaKernel for super::CosKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/cos.ptx"));

diff --git a/src/tensor_ops/cuda_kernels.rs b/src/tensor_ops/cuda_kernels.rs
@@ -3,7 +3,7 @@ use crate::{
     tensor::cuda::{Cuda, CudaArray},
     tensor_ops::ops::{BinaryKernel, UnaryKernel},
 };
-use cudarc::device::{AsKernelParam, CudaSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{AsKernelParam, CudaSlice, LaunchAsync, LaunchConfig};
 use std::sync::Arc;
 
 pub trait UnaryOpCudaKernel {

diff --git a/src/tensor_ops/div/cuda_kernel.rs b/src/tensor_ops/div/cuda_kernel.rs
@@ -1,7 +1,7 @@
 use crate::tensor_ops::cuda_kernels::{BinaryOpCudaKernel, UnaryOpCudaKernel};
 
-unsafe impl cudarc::device::AsKernelParam for super::ScalarDivKernelOp<f32> {}
-unsafe impl cudarc::device::AsKernelParam for super::BinaryDivKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::ScalarDivKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::BinaryDivKernelOp {}
 
 impl UnaryOpCudaKernel for super::ScalarDivKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/scalar_div.ptx"));

diff --git a/src/tensor_ops/exp/cuda_kernel.rs b/src/tensor_ops/exp/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::ExpKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::ExpKernelOp {}
 
 impl UnaryOpCudaKernel for super::ExpKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/exp.ptx"));

diff --git a/src/tensor_ops/huber_error/cuda_kernel.rs b/src/tensor_ops/huber_error/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::BinaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::HuberErrorKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::HuberErrorKernelOp<f32> {}
 
 impl BinaryOpCudaKernel for super::HuberErrorKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/huber_error.ptx"));

diff --git a/src/tensor_ops/ln/cuda_kernel.rs b/src/tensor_ops/ln/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::LnKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::LnKernelOp {}
 
 impl UnaryOpCudaKernel for super::LnKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/ln.ptx"));

diff --git a/src/tensor_ops/maximum/cuda_kernel.rs b/src/tensor_ops/maximum/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::BinaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::MaximumKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::MaximumKernelOp {}
 
 impl BinaryOpCudaKernel for super::MaximumKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/maximum.ptx"));

diff --git a/src/tensor_ops/minimum/cuda_kernel.rs b/src/tensor_ops/minimum/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::BinaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::MinimumKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::MinimumKernelOp {}
 
 impl BinaryOpCudaKernel for super::MinimumKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/minimum.ptx"));

diff --git a/src/tensor_ops/mul/cuda_kernel.rs b/src/tensor_ops/mul/cuda_kernel.rs
@@ -1,7 +1,7 @@
 use crate::tensor_ops::cuda_kernels::{BinaryOpCudaKernel, UnaryOpCudaKernel};
 
-unsafe impl cudarc::device::AsKernelParam for super::ScalarMulKernelOp<f32> {}
-unsafe impl cudarc::device::AsKernelParam for super::BinaryMulKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::ScalarMulKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::BinaryMulKernelOp {}
 
 impl UnaryOpCudaKernel for super::ScalarMulKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/scalar_mul.ptx"));

diff --git a/src/tensor_ops/nans_to/cuda_kernel.rs b/src/tensor_ops/nans_to/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::NansToKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::NansToKernelOp<f32> {}
 
 impl UnaryOpCudaKernel for super::NansToKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/nans_to.ptx"));

diff --git a/src/tensor_ops/negate/cuda_kernel.rs b/src/tensor_ops/negate/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::NegateKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::NegateKernelOp {}
 
 impl UnaryOpCudaKernel for super::NegateKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/negate.ptx"));

diff --git a/src/tensor_ops/permute_to/cuda_kernel.rs b/src/tensor_ops/permute_to/cuda_kernel.rs
@@ -1,7 +1,7 @@
 use crate::shapes::*;
 use crate::tensor::cuda::{Cuda, CudaArray};
 
-use cudarc::device::{LaunchAsync, LaunchConfig};
+use cudarc::driver::{LaunchAsync, LaunchConfig};
 use std::sync::Arc;
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/permute_to.ptx"));

diff --git a/src/tensor_ops/pow/cuda_kernel.rs b/src/tensor_ops/pow/cuda_kernel.rs
@@ -1,7 +1,7 @@
-use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 use super::PowKernelOp;
+use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::PowKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::PowKernelOp<f32> {}
 
 impl UnaryOpCudaKernel for super::PowKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/pow.ptx"));

diff --git a/src/tensor_ops/relu/cuda_kernel.rs b/src/tensor_ops/relu/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::ReLUKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::ReLUKernelOp {}
 
 impl UnaryOpCudaKernel for super::ReLUKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/relu.ptx"));

diff --git a/src/tensor_ops/reshape_to/cuda_kernel.rs b/src/tensor_ops/reshape_to/cuda_kernel.rs
@@ -1,9 +1,8 @@
 use crate::{
-    shapes::{Dtype, HasSameNumelAs, Shape},
+    shapes::{HasSameNumelAs, Shape},
     tensor::cuda::{Cuda, CudaArray},
-    tensor_ops::ops::{BinaryKernel, UnaryKernel},
 };
-use cudarc::device::{AsKernelParam, CudaSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits};
+use cudarc::driver::{CudaSlice, LaunchAsync, LaunchConfig};
 use std::sync::Arc;
 
 const PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/reshape.ptx"));
@@ -45,7 +44,7 @@ impl super::ReshapeKernel<f32> for Cuda {
             &mut storage,      // float *out
             Dst::NUM_DIMS,     // const size_t out_num_dims,
             &dst_dims,         // const size_t *out_dims,
-            &dst_strides       // const size_t *out_strides,
+            &dst_strides,      // const size_t *out_strides,
         );
         unsafe { fwd_fn.launch_async(cfg, params) }?;
 

diff --git a/src/tensor_ops/select_and_gather/cuda_kernel.rs b/src/tensor_ops/select_and_gather/cuda_kernel.rs
@@ -1,11 +1,10 @@
 #![allow(clippy::needless_range_loop)]
 
 use crate::{
-    shapes::{Axes, Dtype, RemoveDimTo, ReplaceDimTo, Shape},
+    shapes::{RemoveDimTo, ReplaceDimTo, Shape},
     tensor::cuda::{Cuda, CudaArray},
-    tensor_ops::ops::{BinaryKernel, UnaryKernel},
 };
-use cudarc::device::{AsKernelParam, CudaSlice, LaunchAsync, LaunchConfig, ValidAsZeroBits};
+use cudarc::driver::{CudaSlice, LaunchAsync, LaunchConfig};
 use std::sync::Arc;
 
 const GATHER_PTX_SRC: &str = include_str!(concat!(env!("OUT_DIR"), "/gather.ptx"));
@@ -24,8 +23,11 @@ impl super::ReplaceDimKernel<f32> for Cuda {
         Src: ReplaceDimTo<Dst, Idx>,
     {
         if !self.dev.has_func(GATHER_MODULE_NAME, GATHER_FWD_FN_NAME) {
-            self.dev
-                .load_ptx(GATHER_PTX_SRC.into(), GATHER_MODULE_NAME, &GATHER_ALL_FN_NAMES)?;
+            self.dev.load_ptx(
+                GATHER_PTX_SRC.into(),
+                GATHER_MODULE_NAME,
+                &GATHER_ALL_FN_NAMES,
+            )?;
         }
 
         let dst = inp.shape.replace(idx.shape);
@@ -37,7 +39,10 @@ impl super::ReplaceDimKernel<f32> for Cuda {
         let inp_strides: CudaSlice<usize> = self.dev.take_async(inp.strides.into())?;
         let idx_strides: CudaSlice<usize> = self.dev.take_async(idx.strides.into())?;
 
-        let fwd_fn = self.dev.get_func(GATHER_MODULE_NAME, GATHER_FWD_FN_NAME).unwrap();
+        let fwd_fn = self
+            .dev
+            .get_func(GATHER_MODULE_NAME, GATHER_FWD_FN_NAME)
+            .unwrap();
         let cfg = LaunchConfig::for_num_elems(numel as u32);
         let params = (
             numel,             // const size_t numel,
@@ -70,15 +75,16 @@ impl super::ReplaceDimKernel<f32> for Cuda {
     where
         Src: ReplaceDimTo<Dst, Idx>,
     {
-        let bwd_fn = self.dev.get_func(GATHER_MODULE_NAME, GATHER_BWD_FN_NAME).unwrap();
+        let bwd_fn = self
+            .dev
+            .get_func(GATHER_MODULE_NAME, GATHER_BWD_FN_NAME)
+            .unwrap();
         let numel = grad_out.data.len();
 
         let inp_dims: CudaSlice<usize> = self.dev.take_async(grad_inp.shape.concrete().into())?;
         let idx_dims: CudaSlice<usize> = self.dev.take_async(idx.shape.concrete().into())?;
-        let out_dims: CudaSlice<usize> = self.dev.take_async(grad_out.shape.concrete().into())?;
         let inp_strides: CudaSlice<usize> = self.dev.take_async(grad_inp.strides.into())?;
         let idx_strides: CudaSlice<usize> = self.dev.take_async(idx.strides.into())?;
-        let out_strides: CudaSlice<usize> = self.dev.take_async(grad_out.strides.into())?;
 
         let cfg = LaunchConfig::for_num_elems(numel as u32);
         let params = (
@@ -115,8 +121,11 @@ impl super::RemoveDimKernel<f32> for Cuda {
         Src: RemoveDimTo<Dst, Idx>,
     {
         if !self.dev.has_func(SELECT_MODULE_NAME, SELECT_FWD_FN_NAME) {
-            self.dev
-                .load_ptx(SELECT_PTX_SRC.into(), SELECT_MODULE_NAME, &SELECT_ALL_FN_NAMES)?;
+            self.dev.load_ptx(
+                SELECT_PTX_SRC.into(),
+                SELECT_MODULE_NAME,
+                &SELECT_ALL_FN_NAMES,
+            )?;
         }
 
         let dst = inp.shape.remove(idx.shape);
@@ -130,7 +139,10 @@ impl super::RemoveDimKernel<f32> for Cuda {
         let idx_strides: CudaSlice<usize> = self.dev.take_async(idx.strides.into())?;
         let dst_strides: CudaSlice<usize> = self.dev.take_async(dst.strides().into())?;
 
-        let fwd_fn = self.dev.get_func(SELECT_MODULE_NAME, SELECT_FWD_FN_NAME).unwrap();
+        let fwd_fn = self
+            .dev
+            .get_func(SELECT_MODULE_NAME, SELECT_FWD_FN_NAME)
+            .unwrap();
         let cfg = LaunchConfig::for_num_elems(numel as u32);
         let params = (
             numel,             // const size_t numel,
@@ -144,7 +156,7 @@ impl super::RemoveDimKernel<f32> for Cuda {
             &idx_strides,      // const size_t *idx_strides,
             &mut storage,      // float *out,
             &dst_dims,         // const size_t *out_dims,
-            &dst_strides       // const size_t *out_strides
+            &dst_strides,      // const size_t *out_strides
         );
         unsafe { fwd_fn.launch_async(cfg, params) }?;
 
@@ -164,7 +176,10 @@ impl super::RemoveDimKernel<f32> for Cuda {
     where
         Src: RemoveDimTo<Dst, Idx>,
     {
-        let bwd_fn = self.dev.get_func(SELECT_MODULE_NAME, SELECT_BWD_FN_NAME).unwrap();
+        let bwd_fn = self
+            .dev
+            .get_func(SELECT_MODULE_NAME, SELECT_BWD_FN_NAME)
+            .unwrap();
         let numel = grad_out.data.len();
 
         let inp_dims: CudaSlice<usize> = self.dev.take_async(grad_inp.shape.concrete().into())?;

diff --git a/src/tensor_ops/sigmoid/cuda_kernel.rs b/src/tensor_ops/sigmoid/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::SigmoidKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::SigmoidKernelOp {}
 
 impl UnaryOpCudaKernel for super::SigmoidKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/sigmoid.ptx"));

diff --git a/src/tensor_ops/sin/cuda_kernel.rs b/src/tensor_ops/sin/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::SinKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::SinKernelOp {}
 
 impl UnaryOpCudaKernel for super::SinKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/sin.ptx"));

diff --git a/src/tensor_ops/sqrt/cuda_kernel.rs b/src/tensor_ops/sqrt/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::SqrtKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::SqrtKernelOp {}
 
 impl UnaryOpCudaKernel for super::SqrtKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/sqrt.ptx"));

diff --git a/src/tensor_ops/square/cuda_kernel.rs b/src/tensor_ops/square/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::SquareKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::SquareKernelOp {}
 
 impl UnaryOpCudaKernel for super::SquareKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/square.ptx"));

diff --git a/src/tensor_ops/sub/cuda_kernel.rs b/src/tensor_ops/sub/cuda_kernel.rs
@@ -1,7 +1,7 @@
 use crate::tensor_ops::cuda_kernels::{BinaryOpCudaKernel, UnaryOpCudaKernel};
 
-unsafe impl cudarc::device::AsKernelParam for super::ScalarSubKernelOp<f32> {}
-unsafe impl cudarc::device::AsKernelParam for super::BinarySubKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::ScalarSubKernelOp<f32> {}
+unsafe impl cudarc::driver::AsKernelParam for super::BinarySubKernelOp {}
 
 impl UnaryOpCudaKernel for super::ScalarSubKernelOp<f32> {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/scalar_sub.ptx"));

diff --git a/src/tensor_ops/sum_to/cuda_kernel.rs b/src/tensor_ops/sum_to/cuda_kernel.rs
@@ -3,7 +3,7 @@ use crate::{
     tensor::cuda::{Cuda, CudaArray},
 };
 
-use cudarc::device::{CudaSlice, LaunchAsync, LaunchConfig};
+use cudarc::driver::{CudaSlice, LaunchAsync, LaunchConfig};
 
 use std::sync::Arc;
 

diff --git a/src/tensor_ops/tanh/cuda_kernel.rs b/src/tensor_ops/tanh/cuda_kernel.rs
@@ -1,6 +1,6 @@
 use crate::tensor_ops::cuda_kernels::UnaryOpCudaKernel;
 
-unsafe impl cudarc::device::AsKernelParam for super::TanhKernelOp {}
+unsafe impl cudarc::driver::AsKernelParam for super::TanhKernelOp {}
 
 impl UnaryOpCudaKernel for super::TanhKernelOp {
     const PTX_SRC: &'static str = include_str!(concat!(env!("OUT_DIR"), "/tanh.ptx"));