torchcell.models.hetero_cell_bipartite_dango.AttentionConvWrapper

class AttentionConvWrapper(conv: Module, target_dim: int, norm: str | None = None, activation: str | None = None, dropout: float = 0.1)

Bases: Module

Wrap a graph conv with projection, normalization, activation, and dropout.

forward(x: Any, edge_index: Tensor, **kwargs: Any) → Tensor

Run the conv, then project, normalize, activate, and drop out.

x is Any: HeteroConv passes either a Tensor or a (src, dst) Tensor tuple per edge type, depending on whether the wrapped conv is bipartite.