enerzyme.models.schnet.interaction.CFConv

class enerzyme.models.schnet.interaction.CFConv(in_channels: int, out_channels: int, num_filters: int, nn: Sequential, cutoff: float)

Bases: MessagePassing

__init__(in_channels: int, out_channels: int, num_filters: int, nn: Sequential, cutoff: float)

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor, edge_index: Tensor, edge_weight: Tensor, edge_attr: Tensor) → Tensor

Runs the forward pass of the module.

message(x_j: Tensor, W: Tensor) → Tensor

Constructs messages from node \(j\) to node \(i\) in analogy to \(\phi_{\mathbf{\Theta}}\) for each edge in edge_index. This function can take any argument as input which was initially passed to propagate(). Furthermore, tensors passed to propagate() can be mapped to the respective nodes \(i\) and \(j\) by appending _i or _j to the variable name, .e.g. x_i and x_j.

reset_parameters()

Resets all learnable parameters of the module.