Spaces:

yuhuili
/

EAGLE

Sleeping

yuhuili commited on Dec 9, 2023

Commit

6e01670

•

1 Parent(s): 9647dd6

Update model/ea_model.py

Files changed (1) hide show

model/ea_model.py CHANGED Viewed

@@ -69,7 +69,7 @@ class EaModel(nn.Module):
         device = base_model.model.layers[-1].self_attn.q_proj.weight.device
-        self.ea_layer.to(self.base_model.dtype).to(device)
         self.ea_layer.init_tree()

         device = base_model.model.layers[-1].self_attn.q_proj.weight.device
+        self.ea_layer.to(torch.float16).to(device)
         self.ea_layer.init_tree()