(self, x, e)
| 469 | self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) |
| 470 | |
| 471 | def forward(self, x, e): |
| 472 | assert e.dtype == torch.float32 |
| 473 | with amp.autocast(dtype=torch.float32, device_type="cuda"): |
| 474 | e = (self.modulation.to(dtype=e.dtype, device=e.device) + e.unsqueeze(1)).chunk(2, dim=1) |
| 475 | x = (self.head(self.norm(x) * (1 + e[1]) + e[0])) |
| 476 | return x |
| 477 | |
| 478 | |
| 479 | class MLPProj(torch.nn.Module): |