diff --git a/ppcls/arch/gears/identity_head.py b/ppcls/arch/gears/identity_head.py
index d4343b9c5499739a1d2af4b9430aae1c1faa1c05..0d30d0e450408464dc201e9ea4fc91fa57f13b92 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__()