(self, in_dim, out_dim, dropout=0.0)
| 186 | class ResidualBlock(nn.Module): |
| 187 | |
| 188 | def __init__(self, in_dim, out_dim, dropout=0.0): |
| 189 | super().__init__() |
| 190 | self.in_dim = in_dim |
| 191 | self.out_dim = out_dim |
| 192 | |
| 193 | # layers |
| 194 | self.residual = nn.Sequential( |
| 195 | RMS_norm(in_dim, images=False), nn.SiLU(), |
| 196 | CausalConv3d(in_dim, out_dim, 3, padding=1), |
| 197 | RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout), |
| 198 | CausalConv3d(out_dim, out_dim, 3, padding=1)) |
| 199 | self.shortcut = CausalConv3d(in_dim, out_dim, 1) \ |
| 200 | if in_dim != out_dim else nn.Identity() |
| 201 | |
| 202 | def forward(self, x, feat_cache=None, feat_idx=[0]): |
| 203 | h = self.shortcut(x) |
nothing calls this directly
no test coverage detected