Update torch_vgpu.py

torch_vgpu.py

@@ -40,36 +40,22 @@ class VGPUDevice:
         self.tensor_cores = None  # Will be initialized when needed
         self.device_name = "vgpu"
         self._register_device()
-        
+
     def _register_device(self):
-        """Register vGPU device using torch.library"""
+        """Register vGPU device using PyTorch's device system"""
         try:
-            # Create library for vGPU backend
-            lib = torch.library.Library(self.device_name, "IMPL")
+            # Register vGPU as a device
+            self._device = torch.device(self.device_name)
             
-            # Register basic tensor operations
-            @torch.library.impl(lib, "aten::empty.memory_format")
-            def empty_impl(size, dtype=None, layout=None, device=None, pin_memory=None, memory_format=None):
-                # Create empty tensor in CPU and move to vGPU
+            # 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)
             
-            @torch.library.impl(lib, "aten::add.Tensor")
-            def add_impl(self, other):
-                # Custom implementation of add operation
-                # Move tensors to CPU, add, then back to vGPU
-                cpu_result = self.cpu() + other.cpu()
-                return to_vgpu(cpu_result, self.vram)
-                
-            @torch.library.impl(lib, "aten::copy_")
-            def copy_impl(self, src, non_blocking=False):
-                # Handle tensor copy operations
-                if not isinstance(src, torch.Tensor):
-                    src = torch.tensor(src)
-                return to_vgpu(src.cpu(), self.vram)
-            
-            # Get device after registration
-            self._device = torch.device(f"{self.device_name}:0")
+            # Register allocator for the device
+            torch.utils.set_default_tensor_type(vgpu_allocator)
                 
         except Exception as e:
             raise RuntimeError(f"Failed to register vGPU device: {str(e)}")
@@ -79,10 +65,10 @@ class VGPUDevice:
         return self.device_name
         
     def __str__(self):
-        return f"{self.device_name}:0"
+        return str(self._device)
         
     def __repr__(self):
-        return f"{self.device_name}:0"
+        return str(self._device)
         
     def device(self):
         """Get the PyTorch device object"""
@@ -90,24 +76,7 @@ class VGPUDevice:
     
     def mode(self):
         """Get a context manager for vGPU operations"""
-        class _VGPUMode:
-            def __init__(self, device):
-                self.device = device
-            
-            def __enter__(self):
-                return self
-                
-            def __exit__(self, exc_type, exc_val, exc_tb):
-                pass
-            
-            def __call__(self, fn):
-                def wrapped(*args, **kwargs):
-                    if 'device' in kwargs:
-                        kwargs['device'] = str(self.device)
-                    return fn(*args, **kwargs)
-                return wrapped
-        
-        return _VGPUMode(self)
+        return torch.device(self._device)
     
     def _init_tensor_cores(self):
         if self.tensor_cores is None: