Method __call__

lwm/vqgan.py:346–351

(self, hidden_states)

    @nn.compact
    def __call__(self, hidden_states):
        hidden_states = ResnetBlock(dropout_prob=self.dropout)(hidden_states)
        if not self.no_attn:
            hidden_states = AttnBlock()(hidden_states)
        hidden_states = ResnetBlock(dropout_prob=self.dropout)(hidden_states)
        return hidden_states