Not ready for review

Summary:
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: b2132ac0f7855738bd1259d19b553bfc16f5d09a

Author: daniellepintz

tests/test_uvm_tensor.py

@@ -20,6 +20,9 @@
 @pytest.mark.cpu_and_gpu
 def test_uvm_tensor() -> None:
     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(