Update model/ea_model.py

model/ea_model.py

@@ -47,8 +47,9 @@ class EaModel(nn.Module):
             bias=con["bias"]
         except:
             bias=True
+        print("draft init")
         self.ea_layer = Model(config,bias=bias,total_tokens=total_token,depth=depth,top_k=top_k,threshold=threshold)
-
+        print("draft init end")
         low_memory=False
 
         device = base_model.model.layers[-1].self_attn.q_proj.weight.device