Fixed tests

predibase · Mar 26, 2024 · 6411587 · 6411587
1 parent c7e6e8e
commit 6411587
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/server/lorax_server/utils/adapter.py b/server/lorax_server/utils/adapter.py
@@ -182,7 +182,7 @@ def merge_adapter_weights(
     Returns:
         Tuple[Dict[str, torch.Tensor], Set[str]]: A tuple containing the merged weights and the set of processed adapter weight names.
     """
-    from lorax_server.adapters.config import LoraConfig
+    from lorax_server.adapters.lora import LoraConfig
 
     if not isinstance(adapter_config, LoraConfig):
         raise ValueError(f"Unsupported adapter config type: {type(adapter_config)}")

diff --git a/server/tests/utils/test_lora.py b/server/tests/utils/test_lora.py
@@ -34,7 +34,7 @@ def test_batched_lora_weights(lora_ranks: List[int]):
         batched_weights.add_adapter(idx, weights)
 
     assert not batched_weights.is_empty()
-    assert len(batched_weights.lora_weights) == 2
+    assert len(batched_weights.adapter_weights) == 2
 
     meta = AdapterBatchMetadata(
         adapter_indices=torch.tensor([0, 0, 1, 1, 0, 0, 1, 1], dtype=torch.int64),