Fix: Handle nn.Module in move_to_device

The move_to_device function was previously only expecting torch.Tensor objects.
This caused an AttributeError when a torch.nn.Module (like a Linear layer)
was passed to it, as modules do not have a `.device` attribute directly
in the same way tensors do.

This commit modifies move_to_device to:
1. Check if the input is an instance of nn.Module.
2. If so, use the module's own `.to(device)` method for device placement.
3. Updates type hints for the function to reflect it can handle
   Union[torch.Tensor, torch.nn.Module].

Additionally, a new unit test has been added to test_api.py to verify
that move_to_device correctly handles both torch.Tensor and torch.nn.Module
objects, moving them to the target device (CPU or CUDA) as expected.
This ensures the fix is effective and prevents regressions.

Author: google-labs-jules[bot]

quantllm/quant/quantization_engine.py

@@ -15,19 +15,25 @@ def get_device_map(model: PreTrainedModel) -> Dict[str, torch.device]:
     return device_map
 
 def move_to_device(
-    tensor: torch.Tensor,
+    tensor: Union[torch.Tensor, torch.nn.Module],
     device: torch.device,
     force_copy: bool = False
-) -> torch.Tensor:
-    """Safely move tensor to device with proper error handling."""
+) -> Union[torch.Tensor, torch.nn.Module]:
+    """Safely move tensor or module to device with proper error handling."""
     try:
+        if isinstance(tensor, torch.nn.Module):
+            return tensor.to(device)
+        # Existing logic for torch.Tensor
         if force_copy:
             return tensor.to(device, copy=True)
-        if tensor.device == device:
+        if tensor.device == device: # type: ignore[union-attr]
             return tensor
         return tensor.to(device)
     except Exception as e:
-        raise RuntimeError(f"Failed to move tensor to {device}: {str(e)}")
+        # It's good practice to indicate which tensor/module failed if possible,
+        # but tensor name isn't available here.
+        type_str = "module" if isinstance(tensor, torch.nn.Module) else "tensor"
+        raise RuntimeError(f"Failed to move {type_str} to {device}: {str(e)}")
 
 class DeviceManager:
     """Manage device placement and synchronization."""

quantllm/quant/tests/test_api.py

@@ -197,5 +197,61 @@ def test_gguf_convert_to_gguf_gpu_layersCpu(self): # Quantizer on GPU, layers on
     def test_gguf_convert_to_gguf_gpu_layersDevice(self): # Quantizer on GPU, layers on GPU
         self._run_gguf_conversion_test(quantizer_device=self.device, gguf_cpu_offload=False)
 
+# New test class for move_to_device
+from quantllm.quant.quantization_engine import move_to_device
+import torch.nn as nn
+
+class TestMoveToDevice(unittest.TestCase):
+    def test_move_tensor_and_module(self):
+        """Test move_to_device with both torch.Tensor and torch.nn.Module."""
+        target_device_str = "cuda" if torch.cuda.is_available() else "cpu"
+        target_device = torch.device(target_device_str)
+
+        # 1. Create a simple torch.Tensor
+        my_tensor = torch.randn(2, 3, device="cpu") # Start on CPU
+
+        # 2. Create a simple torch.nn.Module
+        my_module = nn.Linear(10, 10).to("cpu") # Start on CPU
+
+        # 4. Call move_to_device for the tensor and the module
+        moved_tensor = move_to_device(my_tensor, target_device)
+        moved_module = move_to_device(my_module, target_device)
+
+        # 5. Assert that the tensor is on the target device
+        self.assertEqual(moved_tensor.device, target_device, "Tensor not moved to target device.")
+
+        # 6. Assert that the module is on the target device
+        self.assertIsInstance(moved_module, nn.Module, "move_to_device did not return a Module.")
+        
+        # Check device of a parameter
+        if list(moved_module.parameters()): # Check if module has parameters
+            self.assertEqual(
+                next(moved_module.parameters()).device, 
+                target_device, 
+                "Module's parameters not moved to target device."
+            )
+        else: # Handle modules with no parameters (e.g. nn.ReLU()) if needed for future tests
+            # For a simple Linear layer, this else block shouldn't be hit.
+            # If testing with modules without parameters, one might check an attribute
+            # or skip device check if not applicable. For nn.Linear, parameters exist.
+            pass
+
+        # Test with force_copy=True for tensors
+        another_tensor = torch.randn(2,3, device=target_device)
+        copied_tensor = move_to_device(another_tensor, target_device, force_copy=True)
+        self.assertEqual(copied_tensor.device, target_device)
+        if target_device_str == "cpu": # On CPU, to() without copy=True might return same object if already on device
+             pass # Data pointer check is more complex and not strictly necessary for device check
+        else: # On CUDA, .to(device) typically creates a new tensor unless it's already there.
+             # force_copy=True should ensure it's a different object.
+             if another_tensor.is_cuda and copied_tensor.is_cuda: # Both on CUDA
+                self.assertNotEqual(another_tensor.data_ptr(), copied_tensor.data_ptr(), "force_copy=True did not create a new tensor copy on CUDA.")
+        
+        # Test moving a module already on the target device
+        module_on_target = nn.Linear(5,5).to(target_device)
+        moved_module_again = move_to_device(module_on_target, target_device)
+        self.assertEqual(next(moved_module_again.parameters()).device, target_device)
+
+
 if __name__ == '__main__':
     unittest.main()