Update torch_vgpu.py

torch_vgpu.py

@@ -1,6 +1,3 @@
-"""
-Custom PyTorch device implementation that routes operations through our virtual GPU.
-"""
 import torch
 from torch.library import Library, impl
 from typing import Optional, Union, Tuple
@@ -26,7 +23,7 @@ def init_vgpu_backend():
             
             @impl(impl_lib, "custom_allocate")
             def custom_allocate(device=None):
-                return torch.empty((), device='cpu')
+                return torch.empty((), device="cpu")
             
             @impl(impl_lib, "custom_to_cpu")
             def custom_to_cpu(tensor):
@@ -36,28 +33,31 @@ def init_vgpu_backend():
             def custom_from_cpu(tensor):
                 return tensor.clone()
                 
-            # Register our device type
-            torch._C._register_device_type("vgpu", 10)  # Use a custom type ID
+            # Generate all methods for our backend
+            torch.utils.generate_methods_for_privateuse1_backend(
+                for_tensor=True,
+                for_module=True,
+                for_packed_sequence=True,
+                for_storage=True
+            )
             
-            # Mark initialization as complete
             VGPU_BACKEND_INITIALIZED = True
             
         return VGPU_BACKEND_INITIALIZED
     except Exception as e:
         print(f"Backend initialization warning: {e}")
         return False
-              
 
 class VGPUStorage(torch.Storage):
     """Custom storage class that uses our virtual VRAM"""
     
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.vram = kwargs.get('vram')
+        self.vram = kwargs.get("vram")
         if not self.vram:
             from virtual_vram import VirtualVRAM
             self.vram = VirtualVRAM()
-        self.tensor_id = kwargs.get('tensor_id', f"tensor_{id(self)}")
+        self.tensor_id = kwargs.get("tensor_id", f"tensor_{id(self)}")
     
     def _new_shared(self, size):
         return VGPUStorage(size, vram=self.vram)
@@ -81,7 +81,7 @@ class VGPUDevice:
     def __init__(self, vram: Optional[VirtualVRAM] = None):
         self.vram = vram or VirtualVRAM()
         self.tensor_cores = None  # Will be initialized when needed
-        self.device_name = "vgpu"  # Our registered device type
+        self.device_name = "privateuseone"  # Our registered device type
         self._register_device()
 
     def _register_device(self):
@@ -91,7 +91,7 @@ class VGPUDevice:
                 raise RuntimeError("VGPU backend not properly initialized")
                 
             # Create device using our registered device type
-            self._device = torch.device("vgpu")
+            self._device = torch.device(self.device_name)
             
             # Store this instance for reuse
             VGPUDevice._VGPU_INSTANCES[self.device_name] = self
@@ -187,3 +187,4 @@ def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.T
     # Set the device using the internal name
     result.data = result.data.to(device._device)
     return result
+