(self, x, t, y)
| 225 | return pos_enc |
| 226 | |
| 227 | def forward(self, x, t, y): |
| 228 | t = t.unsqueeze(-1).type(torch.float) |
| 229 | t = self.pos_encoding(t, self.time_dim) |
| 230 | |
| 231 | if y is not None: |
| 232 | t += self.label_emb(y) |
| 233 | |
| 234 | x1 = self.inc(x) |
| 235 | x2 = self.down1(x1, t) |
| 236 | x2 = self.sa1(x2) |
| 237 | x3 = self.down2(x2, t) |
| 238 | x3 = self.sa2(x3) |
| 239 | x4 = self.down3(x3, t) |
| 240 | x4 = self.sa3(x4) |
| 241 | |
| 242 | x4 = self.bot1(x4) |
| 243 | x4 = self.bot2(x4) |
| 244 | x4 = self.bot3(x4) |
| 245 | |
| 246 | x = self.up1(x4, x3, t) |
| 247 | x = self.sa4(x) |
| 248 | x = self.up2(x, x2, t) |
| 249 | x = self.sa5(x) |
| 250 | x = self.up3(x, x1, t) |
| 251 | x = self.sa6(x) |
| 252 | output = self.outc(x) |
| 253 | return output |
| 254 | |
| 255 | |
| 256 | if __name__ == '__main__': |
nothing calls this directly
no test coverage detected