(
self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5
)
| 249 | """ |
| 250 | |
| 251 | def __init__( |
| 252 | self, embedding_dim: int, out_dim: int, num_groups: int, act_fn: Optional[str] = None, eps: float = 1e-5 |
| 253 | ): |
| 254 | super().__init__() |
| 255 | self.num_groups = num_groups |
| 256 | self.eps = eps |
| 257 | |
| 258 | if act_fn is None: |
| 259 | self.act = None |
| 260 | else: |
| 261 | self.act = get_activation(act_fn) |
| 262 | |
| 263 | self.linear = nn.Linear(embedding_dim, out_dim * 2) |
| 264 | |
| 265 | def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor: |
| 266 | if self.act: |
nothing calls this directly
no test coverage detected