torchcell.models.hetero_cell_bipartite_dango_gi.AttentionConvWrapper

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

Bases: Module

Wrap a conv with output projection, optional norm, activation, and dropout.

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

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