r""" Args: x(Tensor): Shape [B, L1, C] e(Tensor): Shape [B, C]
(self, x, e)
| 336 | self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) |
| 337 | |
| 338 | def forward(self, x, e): |
| 339 | r""" |
| 340 | Args: |
| 341 | x(Tensor): Shape [B, L1, C] |
| 342 | e(Tensor): Shape [B, C] |
| 343 | """ |
| 344 | assert e.dtype == torch.float32 |
| 345 | with amp.autocast(dtype=torch.float32): |
| 346 | e = (self.modulation.to(e.device) + e.unsqueeze(1)).chunk(2, dim=1) |
| 347 | x = (self.head(self.norm(x) * (1 + e[1]) + e[0])) |
| 348 | return x |
| 349 | |
| 350 | |
| 351 | class MLPProj(torch.nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected