Update torch_vgpu.py

torch_vgpu.py

@@ -47,56 +47,43 @@ class VGPUDevice:
     def _register_device(self):
         """Register vGPU device using PyTorch's device system"""
         try:
-            # Create internal device using privateuseone backend
+            # Register device using privateuseone
             self._device = torch.device(f"{self.internal_name}:0")
             
-            # Store this instance for device mapping
+            # Store this instance for reuse
             VGPUDevice._VGPU_INSTANCES[self.device_name] = self
             
-            # Register custom dispatcher for device mapping
-            def device_mapper(device_str):
-                if device_str.startswith(self.device_name):
-                    # Map vgpu -> privateuseone
-                    idx = device_str.split(":", 1)[1] if ":" in device_str else "0"
-                    return torch.device(f"{self.internal_name}:{idx}")
-                return None
-            
-            # Register the mapper with PyTorch
-            if not hasattr(torch, '_vgpu_device_mapper'):
-                torch._vgpu_device_mapper = device_mapper
-            
             # Define custom operations for the device
             class VGPUAllocator:
-                def __init__(self, vram):
+                def __init__(self, vram, device):
                     self.vram = vram
+                    self.device = device
                 
                 def __call__(self, size, dtype=None, device=None):
-                    if device is None or str(device).startswith("vgpu"):
-                        # Create tensor on CPU first
-                        cpu_tensor = torch.empty(size, dtype=dtype, device='cpu')
-                        # Move to vGPU storage
-                        return to_vgpu(cpu_tensor, self.vram)
-                    return torch.empty(size, dtype=dtype, device=device)
+                    # Create tensor on CPU first
+                    cpu_tensor = torch.empty(size, dtype=dtype, device='cpu')
+                    # Move to vGPU storage
+                    return to_vgpu(cpu_tensor, self.vram)
             
-            # Set this device as the default for tensor allocation
-            self._allocator = VGPUAllocator(self.vram)
+            # Set up allocator
+            self._allocator = VGPUAllocator(self.vram, self._device)
             
         except Exception as e:
             raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
     
     @property
     def type(self):
-        return self.device_name
+        return self.internal_name
         
     def __str__(self):
-        return f"{self.device_name}:0"
+        return f"{self.internal_name}:0"
         
     def __repr__(self):
-        return f"{self.device_name}:0"
+        return f"vgpu(device='{self.internal_name}:0')"
         
     def device(self):
         """Get the PyTorch device object that maps to our vGPU"""
-        return torch.device(str(self))
+        return self._device  # Return the already created device object
     
     def mode(self):
         """Get a context manager for vGPU operations"""
@@ -154,6 +141,6 @@ def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.T
     result = device._from_vram(tensor_id)
     result.requires_grad = tensor.requires_grad
     
-    # Set the device using the user-facing name
-    result.data = result.data.to(f"{device.device_name}:0")
+    # Set the device using the internal name
+    result.data = result.data.to(device._device)
     return result