(
self,
in_channels: int,
model_dims: int,
encoder_dims: int,
num_heads: int,
num_layers: int = 1,
norm_num_groups: int = 32,
)
| 85 | """A transformer model for inputs with 2 spatial dimensions.""" |
| 86 | |
| 87 | def __init__( |
| 88 | self, |
| 89 | in_channels: int, |
| 90 | model_dims: int, |
| 91 | encoder_dims: int, |
| 92 | num_heads: int, |
| 93 | num_layers: int = 1, |
| 94 | norm_num_groups: int = 32, |
| 95 | ): |
| 96 | super().__init__() |
| 97 | |
| 98 | self.norm = nn.GroupNorm(norm_num_groups, in_channels, pytorch_compatible=True) |
| 99 | self.proj_in = nn.Linear(in_channels, model_dims) |
| 100 | self.transformer_blocks = [ |
| 101 | TransformerBlock(model_dims, num_heads, memory_dims=encoder_dims) |
| 102 | for i in range(num_layers) |
| 103 | ] |
| 104 | self.proj_out = nn.Linear(model_dims, in_channels) |
| 105 | |
| 106 | def __call__(self, x, encoder_x, attn_mask, encoder_attn_mask): |
| 107 | # Save the input to add to the output |
nothing calls this directly
no test coverage detected