(
self,
spatial_dims: int,
dim: int,
num_heads: int,
ffn_expansion_factor: float,
bias: bool,
layer_norm_use_bias: bool = False,
flash_attention: bool = False,
)
| 37 | """ |
| 38 | |
| 39 | def __init__( |
| 40 | self, |
| 41 | spatial_dims: int, |
| 42 | dim: int, |
| 43 | num_heads: int, |
| 44 | ffn_expansion_factor: float, |
| 45 | bias: bool, |
| 46 | layer_norm_use_bias: bool = False, |
| 47 | flash_attention: bool = False, |
| 48 | ): |
| 49 | super().__init__() |
| 50 | self.norm1 = Norm[Norm.INSTANCE, spatial_dims](dim, affine=layer_norm_use_bias) |
| 51 | self.attn = CABlock(spatial_dims, dim, num_heads, bias, flash_attention) |
| 52 | self.norm2 = Norm[Norm.INSTANCE, spatial_dims](dim, affine=layer_norm_use_bias) |
| 53 | self.ffn = FeedForward(spatial_dims, dim, ffn_expansion_factor, bias) |
| 54 | |
| 55 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 56 | x = x + self.attn(self.norm1(x)) |
nothing calls this directly
no test coverage detected