akore commited on
Commit
d10ca53
·
verified ·
1 Parent(s): d583be4

Update modeling_atomformer.py

Browse files
Files changed (1) hide show
  1. modeling_atomformer.py +1 -1
modeling_atomformer.py CHANGED
@@ -2516,7 +2516,7 @@ class AtomformerEncoder(nn.Module):
2516
  for blk in self.blocks:
2517
  input_embeds, pos_embeds = blk(input_embeds, pos_embeds, attention_mask)
2518
 
2519
- return input_embeds, pos_embeds
2520
 
2521
 
2522
  class AtomformerPreTrainedModel(PreTrainedModel): # type: ignore
 
2516
  for blk in self.blocks:
2517
  input_embeds, pos_embeds = blk(input_embeds, pos_embeds, attention_mask)
2518
 
2519
+ return input_embeds[:, :-1]
2520
 
2521
 
2522
  class AtomformerPreTrainedModel(PreTrainedModel): # type: ignore