Any module where forward() takes timestep embeddings as a second argument.
| 63 | |
| 64 | |
| 65 | class TimestepBlock(nn.Module): |
| 66 | """ |
| 67 | Any module where forward() takes timestep embeddings as a second argument. |
| 68 | """ |
| 69 | |
| 70 | @abstractmethod |
| 71 | def forward(self, x, emb): |
| 72 | """ |
| 73 | Apply the module to `x` given `emb` timestep embeddings. |
| 74 | """ |
| 75 | |
| 76 | |
| 77 | class TimestepEmbedSequential(nn.Sequential, TimestepBlock): |
nothing calls this directly
no outgoing calls
no test coverage detected