diff --git a/deepspeed/module_inject/replace_policy.py b/deepspeed/module_inject/replace_policy.py
index 59393b1d8477..e4b1195d9737 100755
--- a/deepspeed/module_inject/replace_policy.py
+++ b/deepspeed/module_inject/replace_policy.py
@@ -372,9 +372,9 @@ def __init__(self, client_module, inference=True, megatron_v2=True):
                 GPTNEOXLayerPolicy._orig_layer_class = None
             else:
                 try:
-                    import megatron
-                    from megatron.model.transformer import ParallelTransformerLayerPipe
-                    GPTNEOXLayerPolicy._orig_layer_class = ParallelTransformerLayerPipe
+                    import transformers
+                    from transformers import GPTNeoXLayer
+                    GPTNEOXLayerPolicy._orig_layer_class = GPTNeoXLayer
                 except ImportError:
                     GPTNEOXLayerPolicy._orig_layer_class = None