enerzyme.models.layers.geometry.RadiusGraphLayer
-
class enerzyme.models.layers.geometry.RadiusGraphLayer(cutoff: float, max_num_neighbors: int)[source]
Bases: BaseFFLayer
-
__init__(cutoff: float, max_num_neighbors: int) → None[source]
-
get_output(Ra: Tensor, batch_seg: Tensor) → Dict[str, Tensor][source]