diff --git a/dietgpu/CMakeLists.txt b/dietgpu/CMakeLists.txt
index 63d8b54..81b4475 100644
--- a/dietgpu/CMakeLists.txt
+++ b/dietgpu/CMakeLists.txt
@@ -9,6 +9,7 @@ add_dependencies(dietgpu
 target_link_libraries(dietgpu PRIVATE
   gpu_float_compress
   glog::glog
+  "${TORCH_LIBRARIES}"
 )
 target_include_directories(dietgpu PRIVATE
   $<BUILD_INTERFACE:${dietgpu_SOURCE_DIR}>
diff --git a/dietgpu/ans_test.py b/dietgpu/ans_test.py
index 2be0f3f..ad84fe7 100644
--- a/dietgpu/ans_test.py
+++ b/dietgpu/ans_test.py
@@ -7,7 +7,7 @@
 
 import torch
 
-torch.ops.load_library("//dietgpu:dietgpu")
+torch.ops.load_library("../build/lib/libdietgpu.so")
 
 
 def run_test(dev, ts, temp_mem=None):
diff --git a/dietgpu/float_test.py b/dietgpu/float_test.py
index 35892b5..ab725e5 100644
--- a/dietgpu/float_test.py
+++ b/dietgpu/float_test.py
@@ -7,7 +7,7 @@
 
 import torch
 
-torch.ops.load_library("//dietgpu:dietgpu")
+torch.ops.load_library("../build/lib/libdietgpu.so")
 
 
 def run_test(dev, ts, temp_mem=None):