(self, x, encoder_x, attn_mask, encoder_attn_mask)
| 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 |
| 108 | input_x = x |
| 109 | dtype = x.dtype |
| 110 | |
| 111 | # Perform the input norm and projection |
| 112 | B, H, W, C = x.shape |
| 113 | x = self.norm(x).reshape(B, -1, C) |
| 114 | x = self.proj_in(x) |
| 115 | |
| 116 | # Apply the transformer |
| 117 | for block in self.transformer_blocks: |
| 118 | x = block(x, encoder_x, attn_mask, encoder_attn_mask) |
| 119 | |
| 120 | # Apply the output projection and reshape |
| 121 | x = self.proj_out(x) |
| 122 | x = x.reshape(B, H, W, C) |
| 123 | |
| 124 | return x + input_x |
| 125 | |
| 126 | |
| 127 | class ResnetBlock2D(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected