(self, x: Tensor, mod: Tensor)
| 127 | super().__init__(*args, **kwargs) |
| 128 | |
| 129 | def forward(self, x: Tensor, mod: Tensor) -> Tensor: |
| 130 | def attn_residual_func(x: Tensor, mod: Tensor) -> Tensor: |
| 131 | return self.ls1(self.attn(self.norm1(x, mod))) |
| 132 | |
| 133 | def ffn_residual_func(x: Tensor, mod: Tensor) -> Tensor: |
| 134 | return self.ls2(self.mlp(self.norm2(x, mod))) |
| 135 | |
| 136 | if self.training and self.sample_drop_ratio > 0.1: |
| 137 | raise NotImplementedError("Modulation with drop path ratio larger than 0.1 is not supported yet") |
| 138 | elif self.training and self.sample_drop_ratio > 0.0: |
| 139 | x = x + self.drop_path1(attn_residual_func(x, mod)) |
| 140 | x = x + self.drop_path1(ffn_residual_func(x, mod)) # FIXME: drop_path2 |
| 141 | else: |
| 142 | x = x + attn_residual_func(x, mod) |
| 143 | x = x + ffn_residual_func(x, mod) |
| 144 | return x |
| 145 | # ******************************************************** |
| 146 | |
| 147 |
nothing calls this directly
no outgoing calls
no test coverage detected