r""" The [`VQModel`] forward method. Args: sample (`torch.Tensor`): Input sample. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`models.autoencoders.vq_model.VQEncoderOutput`] instead of a plain tuple.
(
self, sample: torch.Tensor, return_dict: bool = True
)
| 158 | return DecoderOutput(sample=dec, commit_loss=commit_loss) |
| 159 | |
| 160 | def forward( |
| 161 | self, sample: torch.Tensor, return_dict: bool = True |
| 162 | ) -> Union[DecoderOutput, Tuple[torch.Tensor, ...]]: |
| 163 | r""" |
| 164 | The [`VQModel`] forward method. |
| 165 | |
| 166 | Args: |
| 167 | sample (`torch.Tensor`): Input sample. |
| 168 | return_dict (`bool`, *optional*, defaults to `True`): |
| 169 | Whether or not to return a [`models.autoencoders.vq_model.VQEncoderOutput`] instead of a plain tuple. |
| 170 | |
| 171 | Returns: |
| 172 | [`~models.autoencoders.vq_model.VQEncoderOutput`] or `tuple`: |
| 173 | If return_dict is True, a [`~models.autoencoders.vq_model.VQEncoderOutput`] is returned, otherwise a |
| 174 | plain `tuple` is returned. |
| 175 | """ |
| 176 | |
| 177 | h = self.encode(sample).latents |
| 178 | dec = self.decode(h) |
| 179 | |
| 180 | if not return_dict: |
| 181 | return dec.sample, dec.commit_loss |
| 182 | return dec |