diff --git a/setup.py b/setup.py
index 8a14277f..78db09c9 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 import torch
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CUDA_HOME
 
 __version__ = '0.2.2'
 url = 'https://github.com/rusty1s/pytorch_sparse'
@@ -11,7 +11,7 @@
 ext_modules = []
 cmdclass = {}
 
-if torch.cuda.is_available():
+if CUDA_HOME is not None:
     ext_modules += [
         CUDAExtension(
             'spspmm_cuda',
diff --git a/torch_sparse/spspmm.py b/torch_sparse/spspmm.py
index 6f5d4d3a..9ed5afee 100644
--- a/torch_sparse/spspmm.py
+++ b/torch_sparse/spspmm.py
@@ -37,7 +37,7 @@ def forward(ctx, indexA, valueA, indexB, valueB, m, k, n):
     @staticmethod
     def backward(ctx, grad_indexC, grad_valueC):
         m, k, n = ctx.m, ctx.k, ctx.n
-        indexA, valueA, indexB, valueB, indexC = ctx.saved_variables
+        indexA, valueA, indexB, valueB, indexC = ctx.saved_tensors
 
         grad_valueA = grad_valueB = None