Fix phi (#410)

predibase · Apr 12, 2024 · 4a6a39b · 4a6a39b
1 parent 463ea71
commit 4a6a39b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/server/lorax_server/models/flash_phi.py b/server/lorax_server/models/flash_phi.py
@@ -4,6 +4,7 @@
 import torch.distributed
 from opentelemetry import trace
 from transformers import AutoTokenizer
+from transformers.models.phi.modeling_phi import PhiConfig
 
 from lorax_server.models import FlashCausalLM
 from lorax_server.models.custom_modeling.flash_phi_modeling import (
@@ -14,7 +15,6 @@
     MLP_FC1,
     MLP_FC2,
     FlashPhiForCausalLM,
-    PhiConfig,
 )
 from lorax_server.utils import (
     Weights,