How can I get the coordinates of all the atoms?

[RayTang88]

@lucidrains ,HI,
I read this quote in Readme, set the output atom type by setting the" Atoms" .

Valid choices for atoms include:
backbone - 3 backbone atoms (C, Ca, N) [default]
backbone-with-cbeta - 3 backbone atoms and C beta
backbone-with-oxygen - 3 backbone atoms and oxygen from carboxyl
backbone-with-cbeta-and-oxygen - 3 backbone atoms with C beta and oxygen
all - backbone and all other atoms from sidechain
You can also pass in a tensor of shape (14,) defining which atoms you would like to include

When I set this parameter, I got an error like that,

model = Alphafold2(
    dim = 256,
    depth = 1,
    heads = 8,
    dim_head = 64,
    predict_coords = True,
    # structure_module_dim = 8,
    structure_module_depth = 2,
    structure_module_heads = 4,
    structure_module_dim_head = 16,
    atoms = "all"
    # structure_module_refinement_iters = 2
).to(DEVICE)

Traceback (most recent call last):
  File "/tmp/pycharm_project_152/train_end2end.py", line 99, in <module>
    model = Alphafold2(
TypeError: __init__() got an unexpected keyword argument 'atoms'

Process finished with exit code 1

When I remove this parameter, I can only get the coordinates of one type atom in the sequence in refined.

    refined = model(
            seq,
            msa = msa,
            embedds = embedds,
            mask = mask
        )

Is this atom a CA atom?
#How can I get the coordinates of all the atoms, what parameters should I set, or which version of the previous version can I get ？