fix gpt neo x import

EleutherAI · Oct 15, 2024 · ee37d7b · ee37d7b
1 parent 482eabb
commit ee37d7b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/concept_erasure/scrubbing/neox.py b/concept_erasure/scrubbing/neox.py
@@ -6,10 +6,12 @@
 from tqdm.auto import tqdm
 from transformers import (
     GPTNeoXForCausalLM,
-    GPTNeoXLayer,
     GPTNeoXModel,
 )
-from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXAttention
+from transformers.models.gpt_neox.modeling_gpt_neox import (
+    GPTNeoXAttention,
+    GPTNeoXLayer,
+)
 
 from concept_erasure import ConceptScrubber, ErasureMethod, LeaceFitter
 from concept_erasure.utils import assert_type