Update torch_vgpu.py

torch_vgpu.py

@@ -61,6 +61,33 @@ class VGPUDevice:
         """Get the PyTorch device object"""
         return torch.device(str(self))
     
+    def mode(self):
+        """Get a context manager for vGPU operations"""
+        from torch.utils._python_dispatch import TorchFunctionMode
+        
+        class _VGPUMode(TorchFunctionMode):
+            def __init__(self, device):
+                self.device = device
+            
+            def __torch_function__(self, func, types, args=(), kwargs=None):
+                kwargs = kwargs or {}
+                
+                # Handle tensor creation and device placement
+                if func is torch.tensor or 'device' in kwargs:
+                    kwargs['device'] = str(self.device)
+                
+                # Handle tensor operations
+                new_args = []
+                for arg in args:
+                    if isinstance(arg, torch.Tensor):
+                        if not str(arg.device).startswith(self.device.device_name):
+                            arg = to_vgpu(arg, self.device.vram)
+                    new_args.append(arg)
+                
+                return func(*new_args, **kwargs)
+        
+        return _VGPUMode(self)
+    
     def _init_tensor_cores(self):
         if self.tensor_cores is None:
             from tensor_core import TensorCoreArray
@@ -95,8 +122,11 @@ class VGPUDevice:
         # Create new tensor with result
         return torch.from_numpy(result)
 
-# Register vGPU device type with PyTorch
-torch.backends.register_custom_device("vgpu", VGPUDevice)
+def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
+    """Move a tensor to vGPU device"""
+    device = VGPUDevice(vram)
+    tensor_id = device._to_vram(tensor)
+    return VGPUTensor(device._from_vram(tensor_id))
 
 def to_vgpu(tensor: torch.Tensor, vram: Optional[VirtualVRAM] = None) -> torch.Tensor:
     """Helper function to move tensors to vGPU"""