diff --git a/ppcls/arch/gears/identity_head.py b/ppcls/arch/gears/identity_head.py
index d4343b9c5..0d30d0e45 100644
--- a/ppcls/arch/gears/identity_head.py
+++ b/ppcls/arch/gears/identity_head.py
@@ -1,3 +1,6 @@
+from paddle import nn
+
+
 class IdentityHead(nn.Layer):
     def __init__(self):
         super(IdentityHead, self).__init__()