(self, dim, dim_out, *, time_emb_dim, groups=8)
| 114 | |
| 115 | class ResnetBlock(nn.Module): |
| 116 | def __init__(self, dim, dim_out, *, time_emb_dim, groups=8): |
| 117 | super().__init__() |
| 118 | self.mlp = nn.Sequential( |
| 119 | Mish(), |
| 120 | nn.Linear(time_emb_dim, dim_out) |
| 121 | ) |
| 122 | |
| 123 | self.block1 = Block(dim, dim_out) |
| 124 | self.block2 = Block(dim_out, dim_out) |
| 125 | self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity() |
| 126 | |
| 127 | def forward(self, x, time_emb): |
| 128 | h = self.block1(x) |