(self, dim, dim_cond, single=False)
| 22 | |
| 23 | class CogAdaLayerNorm(torch.nn.Module): |
| 24 | def __init__(self, dim, dim_cond, single=False): |
| 25 | super().__init__() |
| 26 | self.single = single |
| 27 | self.linear = torch.nn.Linear(dim_cond, dim * (2 if single else 6)) |
| 28 | self.norm = torch.nn.LayerNorm(dim, elementwise_affine=True, eps=1e-5) |
| 29 | |
| 30 | |
| 31 | def forward(self, hidden_states, prompt_emb, emb): |