Update torch_vgpu.py

torch_vgpu.py

@@ -7,12 +7,35 @@ from typing import Optional, Union, Tuple
 import numpy as np
 from virtual_vram import VirtualVRAM
 
-# Register and rename privateuse1 backend to vgpu
-try:
-    torch.utils.rename_privateuse1_backend("vgpu")
-except (AttributeError, RuntimeError) as e:
-    # Fallback for older PyTorch versions or if already renamed
-    pass
+# Initialize custom backend
+def init_vgpu_backend():
+    try:
+        # First rename the backend
+        torch.utils.rename_privateuse1_backend("vgpu")
+        
+        # Then generate all the necessary methods
+        torch.utils.generate_methods_for_privateuse1_backend(
+            for_tensor=True,
+            for_module=True,
+            for_packed_sequence=True,
+            for_storage=True
+        )
+        
+        # Register our custom library
+        lib = Library("vgpu", "DEF")
+        lib.define("custom_op(Tensor self) -> Tensor")
+        
+        @impl("vgpu", "custom_op", "Tensor")
+        def custom_op_impl(tensor):
+            return tensor.clone()
+            
+        return True
+    except Exception as e:
+        print(f"Backend initialization warning: {e}")
+        return False
+
+# Initialize the backend
+VGPU_BACKEND_INITIALIZED = init_vgpu_backend()
 
 class VGPUStorage(torch.Storage):
     """Custom storage class that uses our virtual VRAM"""
@@ -47,15 +70,17 @@ class VGPUDevice:
     def __init__(self, vram: Optional[VirtualVRAM] = None):
         self.vram = vram or VirtualVRAM()
         self.tensor_cores = None  # Will be initialized when needed
-        self.internal_name = "privateuse1"  # PyTorch backend name
-        self.device_name = "vgpu"  # User-facing device name
+        self.device_name = "vgpu"  # Both internal and user-facing name
         self._register_device()
 
     def _register_device(self):
         """Register vGPU device using PyTorch's device system"""
         try:
-            # Register device using privateuse1
-            self._device = torch.device(self.internal_name)
+            if not VGPU_BACKEND_INITIALIZED:
+                raise RuntimeError("VGPU backend not properly initialized")
+                
+            # Create device with explicit index
+            self._device = torch.device("vgpu")
             
             # Store this instance for reuse
             VGPUDevice._VGPU_INSTANCES[self.device_name] = self