* `x` has shape `[batch_size, in_channels, height, width]` * `t` has shape `[batch_size, time_channels]`
(self, x: torch.Tensor, t: torch.Tensor)
| 100 | self.dropout = nn.Dropout(dropout) |
| 101 | |
| 102 | def forward(self, x: torch.Tensor, t: torch.Tensor): |
| 103 | """ |
| 104 | * `x` has shape `[batch_size, in_channels, height, width]` |
| 105 | * `t` has shape `[batch_size, time_channels]` |
| 106 | """ |
| 107 | # First convolution layer |
| 108 | h = self.conv1(self.act1(x)) |
| 109 | # Add time embeddings |
| 110 | if self.is_noise: |
| 111 | h += self.time_emb(self.time_act(t))[:, :, None, None] |
| 112 | # Second convolution layer |
| 113 | h = self.conv2(self.dropout(self.act2(h))) |
| 114 | |
| 115 | # Add the shortcut connection and return |
| 116 | return h + self.shortcut(x) |
| 117 | |
| 118 | |
| 119 | class DownBlock(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected