Method forward

pixal3d/models/sparse_structure_vae.py:186–207

(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False)

        self.middle_block.apply(convert_module_to_f32)

    def forward(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False) -> torch.Tensor:
        h = self.input_layer(x)
        h = h.type(self.dtype)

        for block in self.blocks:
            h = block(h)
        h = self.middle_block(h)

        h = h.type(x.dtype)
        h = self.out_layer(h)

        mean, logvar = h.chunk(2, dim=1)

        if sample_posterior:
            std = torch.exp(0.5 * logvar)
            z = mean + std * torch.randn_like(std)
        else:
            z = mean
            
        if return_raw:
            return z, mean, logvar
        return z
        

class SparseStructureDecoder(nn.Module):