(self, x)
| 395 | self.conv = nn.Conv2d(encoder_dims*num_slots, encoder_dims, kernel_size=1, padding=0, stride=1) |
| 396 | |
| 397 | def forward(self, x): |
| 398 | x = einops.rearrange(x, 'b c h w -> b h w c') |
| 399 | x = self.encoder_pos(x) |
| 400 | x = spatial_flatten(x) |
| 401 | x = self.mlp(self.layer_norm(x)) |
| 402 | slots = self.slot_attention(x) |
| 403 | x = spatial_broadcast(slots, (int(self.resolution), int(self.resolution))) |
| 404 | x = self.decoder_pos(x) |
| 405 | x = einops.rearrange(x, 'b n h w c -> b (n c) h w') |
| 406 | out = self.conv(x) |
| 407 | return out |
| 408 | |
| 409 | |
| 410 | class SlotAttentionModule2(nn.Module): |
nothing calls this directly
no test coverage detected