| 1 | import torch |
| 2 | |
| 3 | from ltx_core.model.transformer.gelu_approx import GELUApprox |
| 4 | |
| 5 | |
| 6 | class FeedForward(torch.nn.Module): |
| 7 | def __init__(self, dim: int, dim_out: int, mult: int = 4) -> None: |
| 8 | super().__init__() |
| 9 | inner_dim = int(dim * mult) |
| 10 | project_in = GELUApprox(dim, inner_dim) |
| 11 | |
| 12 | self.net = torch.nn.Sequential(project_in, torch.nn.Identity(), torch.nn.Linear(inner_dim, dim_out)) |
| 13 | |
| 14 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 15 | return self.net(x) |
| 16 |