| 260 | return self.features[idx], self.label_emb[idx], self.label[idx] |
| 261 | |
| 262 | class ProjectHead(nn.Module): |
| 263 | def __init__(self, input_dim: int, target_dim: int) -> None: |
| 264 | super().__init__() |
| 265 | |
| 266 | self.input_dim = input_dim |
| 267 | self.target_dim = target_dim |
| 268 | |
| 269 | self.linear1 = nn.Linear(input_dim, target_dim) |
| 270 | self.linear2 = nn.Linear(target_dim, target_dim) |
| 271 | self.norm = nn.LayerNorm(target_dim) |
| 272 | |
| 273 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 274 | x = self.linear1(x) |
| 275 | x = F.gelu(x) |
| 276 | x = self.linear2(x) |
| 277 | x = self.norm(x) |
| 278 | return x |
| 279 | |
| 280 | class alignMLP(nn.Module): |
| 281 | def __init__(self, input_dim, hidden_dim, output_dim): |
nothing calls this directly
no outgoing calls
no test coverage detected