AI is creating summary for apply_gate Args: x (torch.Tensor): input tensor. gate (torch.Tensor, optional): gate tensor. Defaults to None. tanh (bool, optional): whether to use tanh function. Defaults to False. Returns: torch.Tensor: the output tensor after a
(x, gate=None, tanh=False)
| 168 | |
| 169 | |
| 170 | def apply_gate(x, gate=None, tanh=False): |
| 171 | """AI is creating summary for apply_gate |
| 172 | |
| 173 | Args: |
| 174 | x (torch.Tensor): input tensor. |
| 175 | gate (torch.Tensor, optional): gate tensor. Defaults to None. |
| 176 | tanh (bool, optional): whether to use tanh function. Defaults to False. |
| 177 | |
| 178 | Returns: |
| 179 | torch.Tensor: the output tensor after apply gate. |
| 180 | """ |
| 181 | if gate is None: |
| 182 | return x |
| 183 | if tanh: |
| 184 | return x * gate.unsqueeze(1).tanh() |
| 185 | else: |
| 186 | return x * gate.unsqueeze(1) |
| 187 | |
| 188 | |
| 189 | class RMSNorm(nn.Module): |