Not ready for review (#148)

Summary:
Pull Request resolved: #148

Attempt to fix torchsnapshot CI: https://github.com/pytorch/torchsnapshot/actions/runs/5766115388/job/15694536972
```
tests/test_uvm_tensor.py::test_uvm_tensor FAILED                         [100%]

=================================== FAILURES ===================================
_______________________________ test_uvm_tensor ________________________________

    pytest.mark.cpu_and_gpu
    def test_uvm_tensor() -> None:
        if torch.cuda.is_available() and _UVM_TENSOR_AVAILABLE:
            uvm_tensor = torch.rand(
                (64, 64),
>               out=new_managed_tensor(
                    torch.empty(0, dtype=torch.float32, device="cuda:0"),
                    [64, 64],
                ),
            )

tests/test_uvm_tensor.py:25:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <OpOverloadPacket(op='fbgemm.new_managed_tensor')>
args = (tensor([], device='cuda:0'), [64, 64]), kwargs = {}

    def __call__(self, *args, **kwargs):
        # overloading __call__ to ensure torch.ops.foo.bar()
        # is still callable from JIT
        # We save the function ptr as the `op` attribute on
        # OpOverloadPacket to access it here.
>       return self._op(*args, **kwargs or {})
E       RuntimeError: CUDA error: invalid device ordinal
E       CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
E       For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
E       Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

```

Differential Revision: D48135206

fbshipit-source-id: f3e3006c940026f7cfc5176ed611faba21683faf

Author: daniellepintz

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],
             ),
         )