diff --git a/tests/test_uvm_tensor.py b/tests/test_uvm_tensor.py
index dbd7140..2dfd078 100644
--- a/tests/test_uvm_tensor.py
+++ b/tests/test_uvm_tensor.py
@@ -7,6 +7,8 @@
 
 # pyre-ignore-all-errors[56]
 
+import os
+
 import pytest
 import torch
 from torchsnapshot.uvm_tensor import (
@@ -19,11 +21,16 @@
 
 @pytest.mark.cpu_and_gpu
 def test_uvm_tensor() -> None:
+    os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+    os.environ["TORCH_SHOW_CPP_STACKTRACES"] = "1"
     if torch.cuda.is_available() and _UVM_TENSOR_AVAILABLE:
+        print("_UVM_TENSOR_AVAILABLE", _UVM_TENSOR_AVAILABLE)
+        print("torch.cuda.device_count(): ", torch.cuda.device_count())
+        print("torch.cuda.current_device(): ", torch.cuda.current_device())
         uvm_tensor = torch.rand(
             (64, 64),
             out=new_managed_tensor(
-                torch.empty(0, dtype=torch.float32, device="cuda:0"),
+                torch.empty(0, dtype=torch.float32, device=torch.device("cuda:0")),
                 [64, 64],
             ),
         )