codereview revision

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

Author: brian-dellabetta

src/llmcompressor/modifiers/awq/base.py

@@ -2,7 +2,11 @@
 from typing import Any, Dict, List, Optional, Union
 
 import torch
-from compressed_tensors.utils import align_module_device, update_offload_parameter
+from compressed_tensors.utils import (
+    align_module_device,
+    get_execution_device,
+    update_offload_parameter,
+)
 from loguru import logger
 from pydantic import ConfigDict
 from torch.nn import Module
@@ -590,7 +594,7 @@ def _forward_input_with_kwargs(
         kwargs = input_kwargs or self._module_kwargs
         kwargs = _sanitize_kwargs(kwargs, module)
 
-        inputs = inputs.to(next(module.parameters()).device)
+        inputs = inputs.to(get_execution_device(module))
 
         return module(inputs, **kwargs)[0]