update readme

Author: lucidrains

README.md

@@ -141,8 +141,8 @@ from alphafold3_pytorch import (
 contrived_protein = 'AG'
 
 mock_atompos = [
-    torch.randn(6, 3),   # alanine has 6 non-hydrogen atoms
-    torch.randn(5, 3)    # glycine has 5 non-hydrogen atoms
+    torch.randn(5, 3),   # alanine has 6 non-hydrogen atoms
+    torch.randn(4, 3)    # glycine has 5 non-hydrogen atoms
 ]
 
 train_alphafold3_input = Alphafold3Input(
@@ -193,7 +193,7 @@ batched_eval_atom_input = alphafold3_inputs_to_batched_atom_input(eval_alphafold
 alphafold3.eval()
 sampled_atom_pos = alphafold3(**batched_eval_atom_input.dict())
 
-assert sampled_atom_pos.shape == (1, (6 + 5), 3)
+assert sampled_atom_pos.shape == (1, (5 + 4), 3)
 ```
 
 ## Data preparation