| 209 | |
| 210 | |
| 211 | class PushBlock(torch.nn.Module): |
| 212 | def __init__(self): |
| 213 | super().__init__() |
| 214 | |
| 215 | def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): |
| 216 | res_stack.append(hidden_states) |
| 217 | return hidden_states, time_emb, text_emb, res_stack |
| 218 | |
| 219 | |
| 220 | class PopBlock(torch.nn.Module): |