(self, in_channels=None)
| 216 | |
| 217 | class PopMixBlock(torch.nn.Module): |
| 218 | def __init__(self, in_channels=None): |
| 219 | super().__init__() |
| 220 | self.mix_factor = torch.nn.Parameter(torch.Tensor([0.5])) |
| 221 | self.need_proj = in_channels is not None |
| 222 | if self.need_proj: |
| 223 | self.proj = torch.nn.Linear(in_channels, in_channels) |
| 224 | |
| 225 | def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): |
| 226 | res_hidden_states = res_stack.pop() |