Questions about get_outputs function

[pobo95]

Thank you for your great code!

These days I put some code for my work (calculate atomic virial)
And I have a problem for get_outputs function

I tried to get atomic virial from edge information like this

here is what I fix mace/modules/utils.py
https://github.com/pobo95/mace/blob/5d6c16cd3c01246f8435388aefe8957574cf1a72/mace/modules/utils.py#L76-L135
and
https://github.com/pobo95/mace/blob/5d6c16cd3c01246f8435388aefe8957574cf1a72/mace/modules/utils.py#L176-L219

and 'mace/modules/models.py
https://github.com/pobo95/mace/blob/5d6c16cd3c01246f8435388aefe8957574cf1a72/mace/modules/models.py#L247-L260

but when I run this code
I got TypeError: get_outputs() missing 3 required positional arguments: 'batch', 'edge_index', and 'vectors' this error

here is my questions

1. What is the type of batch, edge_index, vectors? Are these tensor?
2. Is this a problem because I didn't define it?

Thank you for your help again!

[pobo95]

Never mind I solved it :)