Update torch_vgpu.py

torch_vgpu.py

@@ -7,6 +7,13 @@ 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
+
 class VGPUStorage(torch.Storage):
     """Custom storage class that uses our virtual VRAM"""
     
@@ -47,8 +54,8 @@ class VGPUDevice:
     def _register_device(self):
         """Register vGPU device using PyTorch's device system"""
         try:
-            # Register device using privateuseone
-            self._device = torch.device(f"{self.internal_name}:0")
+            # Register device using privateuse1
+            self._device = torch.device(self.internal_name)
             
             # Store this instance for reuse
             VGPUDevice._VGPU_INSTANCES[self.device_name] = self