Update torch_vgpu.py

torch_vgpu.py

@@ -38,37 +38,41 @@ 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"
+        self.device_name = "privateuseone"  # Use privateuseone as base device type
         self._register_device()
 
     def _register_device(self):
         """Register vGPU device using PyTorch's device system"""
         try:
-            # Register vGPU as a device
-            self._device = torch.device(self.device_name)
+            # Create device instance using privateuseone backend
+            self._device = torch.device(f"{self.device_name}:0")
             
-            # Set up default allocator for vGPU device
-            def vgpu_allocator(size, dtype=None, device=None):
-                # 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)
-            
-            # Register allocator for the device
-            torch.utils.set_default_tensor_type(vgpu_allocator)
+            # Define custom operations for the device
+            class VGPUAllocator:
+                def __init__(self, vram):
+                    self.vram = vram
                 
+                def __call__(self, size, dtype=None, device=None):
+                    # 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)
+            
         except Exception as e:
             raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
     
     @property
     def type(self):
-        return self.device_name
+        return "vgpu"  # User-facing device type name
         
     def __str__(self):
-        return str(self._device)
+        return "vgpu"
         
     def __repr__(self):
-        return str(self._device)
+        return "vgpu"
         
     def device(self):
         """Get the PyTorch device object"""
@@ -114,8 +118,14 @@ class VGPUDevice:
 
 def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
     """Move a tensor to vGPU device"""
+    if not isinstance(tensor, torch.Tensor):
+        tensor = torch.tensor(tensor)
+    
     device = VGPUDevice(vram)
     tensor_id = device._to_vram(tensor)
     result = device._from_vram(tensor_id)
     result.requires_grad = tensor.requires_grad
+    
+    # Set the device correctly
+    result.data = result.data.to(device.device())
     return result