diff --git a/build/Makefile b/build/Makefile
index cf3c755..81c04e7 100644
--- a/build/Makefile
+++ b/build/Makefile
@@ -3,7 +3,7 @@ CC = g++
 NVCC = nvcc
 
 # Compiler flags
-CFLAGS = -Wall -Wextra -std=c++11
+CFLAGS = -Wall -Wextra -std=c++11 -I/usr/local/cuda/include
 NVCCFLAGS = -std=c++11
 
 # Directories
@@ -19,7 +19,7 @@ CU_OBJS = $(patsubst $(SRCDIR)/%.cu, $(BUILDDIR)/%.cu.o, $(CU_SRCS))
 
 # CUDA flags and libraries
 CUDAFLAGS = -arch=sm_75
-CUDALIBS = -L/usr/local/cuda-10.2/lib64/ -lcudart -lcuda
+CUDALIBS = -L/usr/local/cuda -lcudart -lcuda
 
 # Rule to build the target
 $(TARGET): $(OBJS) $(CU_OBJS)
diff --git a/build/cuda.cu.o b/build/cuda.cu.o
index 63691b1..c04ec68 100644
Binary files a/build/cuda.cu.o and b/build/cuda.cu.o differ
diff --git a/build/libtensor.so b/build/libtensor.so
index 80dc3bc..69817f2 100755
Binary files a/build/libtensor.so and b/build/libtensor.so differ
diff --git a/build/tensor.o b/build/tensor.o
index 377f3e8..3b6337d 100644
Binary files a/build/tensor.o and b/build/tensor.o differ
diff --git a/csrc/tensor.cpp b/csrc/tensor.cpp
index 6e2e336..ad3f4f1 100644
--- a/csrc/tensor.cpp
+++ b/csrc/tensor.cpp
@@ -100,7 +100,8 @@ extern "C" {
                 tensor->device = device;
             }
         #else
-            printf("Warning: CUDA is not available. Cannot perform GPU operations.\n");
+            fprintf(stderr, "ERROR: CUDA is not available. Cannot perform GPU operations.\n");
+            exit(0);
         #endif
     }
 
diff --git a/norch/__pycache__/tensor.cpython-38.pyc b/norch/__pycache__/tensor.cpython-38.pyc
index 2068305..84e4a97 100644
Binary files a/norch/__pycache__/tensor.cpython-38.pyc and b/norch/__pycache__/tensor.cpython-38.pyc differ
diff --git a/norch/tensor.py b/norch/tensor.py
index 779c1c3..66a1a59 100644
--- a/norch/tensor.py
+++ b/norch/tensor.py
@@ -206,6 +206,7 @@ class Tensor:
         result_data = Tensor()
         result_data.tensor = result_tensor_ptr
         result_data.shape = self.shape.copy()
+        result_data.device = self.device
         result_data.ndim = 2
 
         return result_data