(
self, h: torch.Tensor, force_not_quantize: bool = False, return_dict: bool = True, shape=None
)
| 138 | |
| 139 | @apply_forward_hook |
| 140 | def decode( |
| 141 | self, h: torch.Tensor, force_not_quantize: bool = False, return_dict: bool = True, shape=None |
| 142 | ) -> Union[DecoderOutput, torch.Tensor]: |
| 143 | # also go through quantization layer |
| 144 | if not force_not_quantize: |
| 145 | quant, commit_loss, _ = self.quantize(h) |
| 146 | elif self.config.lookup_from_codebook: |
| 147 | quant = self.quantize.get_codebook_entry(h, shape) |
| 148 | commit_loss = torch.zeros((h.shape[0])).to(h.device, dtype=h.dtype) |
| 149 | else: |
| 150 | quant = h |
| 151 | commit_loss = torch.zeros((h.shape[0])).to(h.device, dtype=h.dtype) |
| 152 | quant2 = self.post_quant_conv(quant) |
| 153 | dec = self.decoder(quant2, quant if self.config.norm_type == "spatial" else None) |
| 154 | |
| 155 | if not return_dict: |
| 156 | return dec, commit_loss |
| 157 | |
| 158 | return DecoderOutput(sample=dec, commit_loss=commit_loss) |
| 159 | |
| 160 | def forward( |
| 161 | self, sample: torch.Tensor, return_dict: bool = True |
no test coverage detected