Update modeling_atomformer.py

modeling_atomformer.py

@@ -2516,7 +2516,7 @@ class AtomformerEncoder(nn.Module):
         for blk in self.blocks:
             input_embeds, pos_embeds = blk(input_embeds, pos_embeds, attention_mask)
 
-        return input_embeds[:, :-1]
+        return input_embeds, pos_embeds
 
 
 class AtomformerPreTrainedModel(PreTrainedModel):  # type: ignore
@@ -2550,7 +2550,7 @@ class AtomformerModel(AtomformerPreTrainedModel):
     ) -> torch.Tensor:
         """Forward function call for the transformer model."""
         output: torch.Tensor = self.encoder(input_ids, coords, attention_mask)
-        return output
+        return output[0][:, :-1]
 
 
 class AtomformerForMaskedAM(AtomformerPreTrainedModel):