SelfAttention2d

class torchelie.nn.SelfAttention2d(ch: int, num_heads: int = 1, out_ch: Optional[int] = None)

Self Attention such as used in SAGAN or BigGAN.

Parameters

ch (int) – number of input / output channels

forward(x: torch.Tensor) → torch.Tensor

training: bool