diff --git a/cuda/lltm_cuda.cpp b/cuda/lltm_cuda.cpp
index 2434776..be76798 100644
--- a/cuda/lltm_cuda.cpp
+++ b/cuda/lltm_cuda.cpp
@@ -25,8 +25,8 @@ std::vector<torch::Tensor> lltm_cuda_backward(
 // C++ interface
 
 // NOTE: AT_ASSERT has become AT_CHECK on master after 0.4.
-#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
-#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
 #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
 
 std::vector<torch::Tensor> lltm_forward(
diff --git a/cuda/lltm_cuda_kernel.cu b/cuda/lltm_cuda_kernel.cu
index 02bb9ad..9d6acbb 100644
--- a/cuda/lltm_cuda_kernel.cu
+++ b/cuda/lltm_cuda_kernel.cu
@@ -25,7 +25,7 @@ __device__ __forceinline__ scalar_t d_tanh(scalar_t z) {
 
 template <typename scalar_t>
 __device__ __forceinline__ scalar_t elu(scalar_t z, scalar_t alpha = 1.0) {
-  return fmaxf(0.0, z) + fminf(0.0, alpha * (exp(z) - 1.0));
+  return fmax((scalar_t) 0.0, z) + fmin((scalar_t) 0.0, alpha * (exp(z) - 1.0));
 }
 
 template <typename scalar_t>
@@ -37,13 +37,13 @@ __device__ __forceinline__ scalar_t d_elu(scalar_t z, scalar_t alpha = 1.0) {
 
 template <typename scalar_t>
 __global__ void lltm_cuda_forward_kernel(
-    const torch::PackedTensorAccessor<scalar_t,3,torch::RestrictPtrTraits,size_t> gates,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> old_cell,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> new_h,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> new_cell,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> input_gate,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> output_gate,
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> candidate_cell) {
+    const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> gates,
+    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> old_cell,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> new_h,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> new_cell,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> input_gate,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> output_gate,
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> candidate_cell) {
   //batch index
   const int n = blockIdx.y;
   // column index
@@ -60,15 +60,15 @@ __global__ void lltm_cuda_forward_kernel(
 
 template <typename scalar_t>
 __global__ void lltm_cuda_backward_kernel(
-    torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> d_old_cell,
-    torch::PackedTensorAccessor<scalar_t,3,torch::RestrictPtrTraits,size_t> d_gates,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> grad_h,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> grad_cell,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> new_cell,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> input_gate,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> output_gate,
-    const torch::PackedTensorAccessor<scalar_t,2,torch::RestrictPtrTraits,size_t> candidate_cell,
-    const torch::PackedTensorAccessor<scalar_t,3,torch::RestrictPtrTraits,size_t> gate_weights) {
+    torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> d_old_cell,
+    torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> d_gates,
+    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> grad_h,
+    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> grad_cell,
+    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> new_cell,
+    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> input_gate,
+    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> output_gate,
+    const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> candidate_cell,
+    const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> gate_weights) {
   //batch index
   const int n = blockIdx.y;
   // column index
@@ -116,15 +116,15 @@ std::vector<torch::Tensor> lltm_cuda_forward(
   const int threads = 1024;
   const dim3 blocks((state_size + threads - 1) / threads, batch_size);
 
-  AT_DISPATCH_FLOATING_TYPES(gates.type(), "lltm_forward_cuda", ([&] {
+  AT_DISPATCH_FLOATING_TYPES(gates.scalar_type(), "lltm_forward_cuda", ([&] {
     lltm_cuda_forward_kernel<scalar_t><<<blocks, threads>>>(
-        gates.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),
-        old_cell.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        new_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        new_cell.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        input_gate.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        output_gate.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        candidate_cell.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>());
+        gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),
+        old_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        new_h.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        new_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        input_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        output_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        candidate_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
   }));
 
   return {new_h, new_cell, input_gate, output_gate, candidate_cell, X, gates};
@@ -149,17 +149,17 @@ std::vector<torch::Tensor> lltm_cuda_backward(
   const int threads = 1024;
   const dim3 blocks((state_size + threads - 1) / threads, batch_size);
 
-  AT_DISPATCH_FLOATING_TYPES(X.type(), "lltm_forward_cuda", ([&] {
+  AT_DISPATCH_FLOATING_TYPES(X.scalar_type(), "lltm_forward_cuda", ([&] {
     lltm_cuda_backward_kernel<scalar_t><<<blocks, threads>>>(
-        d_old_cell.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        d_gates.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>(),
-        grad_h.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        grad_cell.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        new_cell.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        input_gate.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        output_gate.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        candidate_cell.packed_accessor<scalar_t,2,torch::RestrictPtrTraits,size_t>(),
-        gates.packed_accessor<scalar_t,3,torch::RestrictPtrTraits,size_t>());
+        d_old_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        d_gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),
+        grad_h.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        grad_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        new_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        input_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        output_gate.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        candidate_cell.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),
+        gates.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>());
   }));
 
   auto d_gate_weights = d_gates.flatten(1, 2);