diff --git a/.travis.yml b/.travis.yml
index ccddbe1c..531e1c0e 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -10,35 +10,35 @@ env:
     - CUDA_HOME=/usr/local/cuda
   jobs:
     # Python 3.6
-    # - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cpu
-    # - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cu92
-    # - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cu101
+    - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cpu
+    - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cu92
+    - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cu101
     - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cu102
     - PYTHON_VERSION=3.6 TORCH_VERSION=1.7.0 IDX=cu110
-    # - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cpu
-    # - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cu92
-    # - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cu101
-    # - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cu102
+    - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cpu
+    - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cu92
+    - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cu101
+    - PYTHON_VERSION=3.6 TORCH_VERSION=1.6.0 IDX=cu102
     # Python 3.7
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cpu
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu92
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu101
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu102
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu110
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cpu
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cu92
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cu101
-    # - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cu102
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cpu
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu92
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu101
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu102
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.7.0 IDX=cu110
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cpu
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cu92
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cu101
+    - PYTHON_VERSION=3.7 TORCH_VERSION=1.6.0 IDX=cu102
     # Python 3.8
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cpu
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu92
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu101
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu102
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu110
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cpu
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cu92
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cu101
-    # - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cu102
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cpu
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu92
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu101
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu102
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.7.0 IDX=cu110
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cpu
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cu92
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cu101
+    - PYTHON_VERSION=3.8 TORCH_VERSION=1.6.0 IDX=cu102
 
 jobs:
   # For daily builds, we only test on Python 3.6 with CUDA 10.2/11.0.
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 87a4f231..9fdaf7a4 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,7 +1,7 @@
 cmake_minimum_required(VERSION 3.0)
 project(torchcluster)
 set(CMAKE_CXX_STANDARD 14)
-set(TORCHCLUSTER_VERSION 1.6.0)
+set(TORCHCLUSTER_VERSION 1.5.8)
 
 option(WITH_CUDA "Enable CUDA support" OFF)
 
diff --git a/script/cuda.sh b/script/cuda.sh
index 2c59ac41..2dba421c 100755
--- a/script/cuda.sh
+++ b/script/cuda.sh
@@ -99,21 +99,12 @@ if [ "${TRAVIS_OS_NAME}" = "linux" ] && [ "${IDX}" = "cu110" ]; then
   sudo dpkg -i cuda-repo-ubuntu1804-11-0-local_11.0.3-450.51.06-1_amd64.deb
   sudo apt-key add /var/cuda-repo-ubuntu1804-11-0-local/7fa2af80.pub
   sudo apt update -qq
-  sudo apt search cuda
-  sudo apt install cuda-nvcc-11-0
-  sudo apt install cuda-libraries-dev-11-0
+  sudo apt install cuda-nvcc-11-0 cuda-libraries-dev-11-0
   sudo apt clean
   CUDA_HOME=/usr/local/cuda-${CUDA_SHORT}
   LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
   PATH=${CUDA_HOME}/bin:${PATH}
   nvcc --version
-
-  echo "LIB"
-  ls $CUDA_HOME/lib
-  echo "LIB64"
-  ls $CUDA_HOME/lib64
-  echo "LIB"
-  ls $CUDA_HOME/include
 fi
 
 if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "${IDX}" != "cpu" ]; then
@@ -122,7 +113,7 @@ if [ "${TRAVIS_OS_NAME}" = "windows" ] && [ "${IDX}" != "cpu" ]; then
   curl -k -L "https://drive.google.com/u/0/uc?id=1injUyo3lnarMgWyRcXqKg4UGnN0ysmuq&export=download" --output "/tmp/gpu_driver_dlls.zip"
   7z x "/tmp/gpu_driver_dlls.zip" -o"/c/Windows/System32"
 
-  # Install CUDA
+  # Install CUDA:
   wget -nv "${CUDA_URL}/${CUDA_FILE}"
   PowerShell -Command "Start-Process -FilePath \"${CUDA_FILE}\" -ArgumentList \"-s nvcc_${CUDA_SHORT} cuobjdump_${CUDA_SHORT} nvprune_${CUDA_SHORT} cupti_${CUDA_SHORT} cublas_dev_${CUDA_SHORT} cudart_${CUDA_SHORT} cufft_dev_${CUDA_SHORT} curand_dev_${CUDA_SHORT} cusolver_dev_${CUDA_SHORT} cusparse_dev_${CUDA_SHORT} npp_dev_${CUDA_SHORT} nvrtc_dev_${CUDA_SHORT} nvml_dev_${CUDA_SHORT}\" -Wait -NoNewWindow"
   CUDA_HOME=/c/Program\ Files/NVIDIA\ GPU\ Computing\ Toolkit/CUDA/v${CUDA_SHORT}
diff --git a/setup.py b/setup.py
index ce6cabb6..9fa90e1a 100644
--- a/setup.py
+++ b/setup.py
@@ -77,7 +77,7 @@ def get_extensions():
 
 setup(
     name='torch_cluster',
-    version='1.6.0',
+    version='1.5.8',
     author='Matthias Fey',
     author_email='matthias.fey@tu-dortmund.de',
     url='https://github.com/rusty1s/pytorch_cluster',
diff --git a/torch_cluster/__init__.py b/torch_cluster/__init__.py
index ea8093e4..4a2ae57d 100644
--- a/torch_cluster/__init__.py
+++ b/torch_cluster/__init__.py
@@ -3,7 +3,7 @@
 
 import torch
 
-__version__ = '1.6.0'
+__version__ = '1.5.8'
 
 for library in [
         '_version', '_grid', '_graclus', '_fps', '_rw', '_sampler', '_nearest',