❓ [QUESTION] About the data class AtomicData about nequip HOT 3 CLOSED

Hi @floatingCatty ,

Thanks for your interest in our code!

Good question: we use the dict type in the model for compatibility with TorchScript and compilation for deployment. I think the most straightforward way to achieve what you want is to keep around the original Batch object that you use as the input to the model, and to copy fields from the output dict of the model back into the input Batch object on which you can call get_example. (This of course assumes that you are mostly just predicting some new per-node/per-edge/per-graph quantities, and not changing the graph structure in the model, but of course if you are changing the graph structure this is a very different scenario.)

from nequip.

Thanks !

I will try to implement what you said. I just want to predict some edge features, so what you just suggest is applicable.

from nequip.

Sounds good @floatingCatty ! For that, the easiest thing may be to look at the --output and --output-fields options of nequip-evaluate as well.

from nequip.