| 149 | |
| 150 | |
| 151 | class ValueFunctionMidBlock1D(nn.Module): |
| 152 | def __init__(self, in_channels: int, out_channels: int, embed_dim: int): |
| 153 | super().__init__() |
| 154 | self.in_channels = in_channels |
| 155 | self.out_channels = out_channels |
| 156 | self.embed_dim = embed_dim |
| 157 | |
| 158 | self.res1 = ResidualTemporalBlock1D(in_channels, in_channels // 2, embed_dim=embed_dim) |
| 159 | self.down1 = Downsample1D(out_channels // 2, use_conv=True) |
| 160 | self.res2 = ResidualTemporalBlock1D(in_channels // 2, in_channels // 4, embed_dim=embed_dim) |
| 161 | self.down2 = Downsample1D(out_channels // 4, use_conv=True) |
| 162 | |
| 163 | def forward(self, x: torch.Tensor, temb: torch.Tensor | None = None) -> torch.Tensor: |
| 164 | x = self.res1(x, temb) |
| 165 | x = self.down1(x) |
| 166 | x = self.res2(x, temb) |
| 167 | x = self.down2(x) |
| 168 | return x |
| 169 | |
| 170 | |
| 171 | class MidResTemporalBlock1D(nn.Module): |
no outgoing calls
no test coverage detected
searching dependent graphs…