Fix Lora scaling device issue

Signed-off-by: Kevin Hsieh <[email protected]>

Author: Hsieh, Kevin

TrainingExtensions/torch/src/python/aimet_torch/peft.py

@@ -71,12 +71,14 @@ def __init__(self, lora_layer: PeftLoraLayer):
         self.base_layer = lora_layer.base_layer
         self.r = lora_layer.r
         self.lora_alpha = lora_layer.lora_alpha
-        self.scaling = [
-            torch.nn.Parameter(torch.as_tensor(scale), requires_grad=False).to(
-                self.base_layer.weight.device
-            )
-            for scale in lora_layer.scaling.values()
-        ]
+        self.scaling = torch.nn.ParameterList(
+            [
+                torch.nn.Parameter(torch.as_tensor(scale), requires_grad=False).to(
+                    self.base_layer.weight.device
+                )
+                for scale in lora_layer.scaling.values()
+            ]
+        )
         self.lora_dropout = nn.ModuleList({})
         self.adapter_name_to_index = {}
         self.index_to_adapter_name = {}

TrainingExtensions/torch/test/python/test_peft.py

@@ -380,10 +380,23 @@ def forward_pass(model, forward_pass_callback=None):
             ]
             assert sorted(tensor_name) == sorted(tensors)
 
+    @pytest.mark.cuda
+    def test_changing_lora_device(self):
+        model = one_adapter_model().cuda()
+
+        replace_lora_layers_with_quantizable_layers(model)
+        dummy_inputs = torch.randn(10, 10).cuda()
+
+        _ = model(dummy_inputs)
+
+        model.cpu()
+
+        _ = model(dummy_inputs.cpu())
+
 
 def _is_frozen(quantizer):
     return (
-        quantizer._allow_overwrite == False
-        and quantizer.min.requires_grad == False
-        and quantizer.max.requires_grad == False
+        not quantizer._allow_overwrite
+        and not quantizer.min.requires_grad
+        and not quantizer.max.requires_grad
     )