From 3e975bb841cf6fcb1b4832bed74d894de1b4ffcc Mon Sep 17 00:00:00 2001 From: Jedrzej Kosinski Date: Tue, 2 Apr 2024 23:40:48 -0500 Subject: [PATCH] Fix for cases where predefined weights length doesn't match expected layer lengths (will need to generalize for for ControlLora to properly do soft weights for actual layer number) --- adv_control/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/adv_control/utils.py b/adv_control/utils.py index e2fd229..c96cb3f 100644 --- a/adv_control/utils.py +++ b/adv_control/utils.py @@ -53,9 +53,12 @@ class ControlWeights: self.weights.reverse() self.weight_mask = weight_mask - def get(self, idx: int) -> Union[float, Tensor]: + def get(self, idx: int, default=1.0) -> Union[float, Tensor]: # if weights is not none, return index if self.weights is not None: + # this implies weights list is not aligning with expectations - will need to adjust code + if idx >= len(self.weights): + return default return self.weights[idx] return 1.0 @@ -542,6 +545,7 @@ class AdvancedControlBase: self.prepare_weight_mask_cond_hint(x, self.batched_number) # adjust mask for current layer and return return torch.pow(self.weight_mask_cond_hint, self.get_calc_pow(idx=idx, layers=layers)) + return self.weights.get(idx=idx) def get_calc_pow(self, idx: int, layers: int) -> int: