(self, x: Tensor)
| 84 | self.sample_drop_ratio = drop_path |
| 85 | |
| 86 | def forward(self, x: Tensor) -> Tensor: |
| 87 | def attn_residual_func(x: Tensor) -> Tensor: |
| 88 | return self.ls1(self.attn(self.norm1(x))) |
| 89 | |
| 90 | def ffn_residual_func(x: Tensor) -> Tensor: |
| 91 | return self.ls2(self.mlp(self.norm2(x))) |
| 92 | |
| 93 | if self.training and self.sample_drop_ratio > 0.1: |
| 94 | # the overhead is compensated only for a drop path rate larger than 0.1 |
| 95 | x = drop_add_residual_stochastic_depth( |
| 96 | x, |
| 97 | residual_func=attn_residual_func, |
| 98 | sample_drop_ratio=self.sample_drop_ratio, |
| 99 | ) |
| 100 | x = drop_add_residual_stochastic_depth( |
| 101 | x, |
| 102 | residual_func=ffn_residual_func, |
| 103 | sample_drop_ratio=self.sample_drop_ratio, |
| 104 | ) |
| 105 | elif self.training and self.sample_drop_ratio > 0.0: |
| 106 | x = x + self.drop_path1(attn_residual_func(x)) |
| 107 | x = x + self.drop_path1(ffn_residual_func(x)) # FIXME: drop_path2 |
| 108 | else: |
| 109 | x = x + attn_residual_func(x) |
| 110 | x = x + ffn_residual_func(x) |
| 111 | return x |
| 112 | |
| 113 | |
| 114 | def drop_add_residual_stochastic_depth( |
no test coverage detected