fix a bug

AgentMaker · Mar 31, 2021 · 09e779f · 09e779f
1 parent 2f80d94
commit 09e779f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ppim/models/deit.py b/ppim/models/deit.py
@@ -68,10 +68,10 @@ def __init__(self,
 
         if class_dim > 0:
             self.head_dist = nn.Linear(self.embed_dim, self.class_dim)
+            self.head_dist.apply(self._init_weights)
 
         trunc_normal_(self.dist_token)
         trunc_normal_(self.pos_embed)
-        self.head_dist.apply(self._init_weights)
 
     def forward_features(self, x):
         B = paddle.shape(x)[0]