(self, x: Tensor)
| 93 | self.sample_drop_ratio = drop_path |
| 94 | |
| 95 | def forward(self, x: Tensor) -> Tensor: |
| 96 | def attn_residual_func(x: Tensor) -> Tensor: |
| 97 | return self.ls1(self.attn(self.norm1(x))) |
| 98 | |
| 99 | def ffn_residual_func(x: Tensor) -> Tensor: |
| 100 | return self.ls2(self.mlp(self.norm2(x))) |
| 101 | |
| 102 | if self.training and self.sample_drop_ratio > 0.1: |
| 103 | # the overhead is compensated only for a drop path rate larger than 0.1 |
| 104 | x = drop_add_residual_stochastic_depth( |
| 105 | x, |
| 106 | residual_func=attn_residual_func, |
| 107 | sample_drop_ratio=self.sample_drop_ratio, |
| 108 | ) |
| 109 | x = drop_add_residual_stochastic_depth( |
| 110 | x, |
| 111 | residual_func=ffn_residual_func, |
| 112 | sample_drop_ratio=self.sample_drop_ratio, |
| 113 | ) |
| 114 | elif self.training and self.sample_drop_ratio > 0.0: |
| 115 | x = x + self.drop_path1(attn_residual_func(x)) |
| 116 | x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 |
| 117 | else: |
| 118 | x = x + attn_residual_func(x) |
| 119 | x = x + ffn_residual_func(x) |
| 120 | return x |
| 121 | |
| 122 | |
| 123 | # ********** Modified by Zexin He in 2023-2024 ********** |
no test coverage detected