(self, cfg: Config, *, parent: Optional[Module])
| 573 | ) |
| 574 | |
| 575 | def __init__(self, cfg: Config, *, parent: Optional[Module]): |
| 576 | super().__init__(cfg, parent=parent) |
| 577 | cfg = self.config |
| 578 | self._add_child( |
| 579 | "attention", cfg.attention.set(target_dim=cfg.input_dim, source_dim=cfg.input_dim) |
| 580 | ) |
| 581 | self._add_child("feed_forward", cfg.feed_forward.set(input_dim=cfg.input_dim)) |
| 582 | self._add_child("adaln", cfg.adaln.set(dim=cfg.input_dim)) |
| 583 | |
| 584 | def forward( |
| 585 | self, |