Update torch_vgpu.py

torch_vgpu.py

@@ -15,11 +15,8 @@ def init_vgpu_backend():
     global VGPU_BACKEND_INITIALIZED
     try:
         if not VGPU_BACKEND_INITIALIZED:
-            # First rename privateuse1 to vgpu
-            torch.utils.rename_privateuse1_backend("vgpu")
-            
-            # Create an aten library implementation for our backend
-            lib = Library("aten", "IMPL", "vgpu")
+            # Create an aten library implementation for our backend with PrivateUse1 dispatch key
+            lib = Library("aten", "IMPL", "PrivateUse1")
             
             # Register essential operations for the backend
             @impl(lib, "empty.memory_format")
@@ -35,9 +32,12 @@ def init_vgpu_backend():
             @impl(lib, "_to_copy")
             def to_impl(self, dtype=None, device=None, non_blocking=False, copy=False):
                 # Handle tensor device transfer
-                if device and str(device).startswith("vgpu"):
+                if device and str(device).startswith("privateuse1"):
                     return self.cpu().clone()
                 return self.clone()
+                
+            # Finally rename privateuse1 to vgpu after implementing the ops
+            torch.utils.rename_privateuse1_backend("vgpu")
             
             # Generate all methods for our backend
             torch.utils.generate_methods_for_privateuse1_backend(
@@ -97,8 +97,8 @@ class VGPUDevice:
             if not VGPU_BACKEND_INITIALIZED:
                 raise RuntimeError("VGPU backend not properly initialized")
                 
-            # Create device with explicit index
-            self._device = torch.device("vgpu:0")
+            # Create device with explicit index - use privateuse1 since that's the actual device type
+            self._device = torch.device("privateuse1:0")
             
             # Store this instance for reuse
             VGPUDevice._VGPU_INSTANCES[self.device_name] = self