enerzyme.models.layers.gather_embedding.GatherAtomEmbedding#

class enerzyme.models.layers.gather_embedding.GatherAtomEmbedding[source]#

Bases: BaseFFLayer

__init__() → None[source]#

get_output(**relevant_input: Dict[str, Tensor]) → Dict[str, Tensor][source]#

get_relevant_input_fields(net_input_fields: Set[str]) → Set[str][source]#