| 126 | return t_emb |
| 127 | |
| 128 | class SequenceEmbed(nn.Module): |
| 129 | def __init__( |
| 130 | self, |
| 131 | dim_w, |
| 132 | dim_h, |
| 133 | norm_layer=None, |
| 134 | bias=True, |
| 135 | ): |
| 136 | super().__init__() |
| 137 | |
| 138 | self.proj = nn.Linear(dim_w, dim_h, bias=bias) |
| 139 | self.norm = norm_layer(dim_h) if norm_layer else nn.Identity() |
| 140 | |
| 141 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 142 | return self.norm(self.proj(x)) |
| 143 | |
| 144 | |
| 145 | class FMTBlock(nn.Module): |