(self, input, return_pred_indices=False)
| 236 | return dec |
| 237 | |
| 238 | def forward(self, input, return_pred_indices=False): |
| 239 | quant, diff, (_,_,ind) = self.encode(input) |
| 240 | dec = self.decode(quant) |
| 241 | if return_pred_indices: |
| 242 | return dec, diff, ind |
| 243 | return dec, diff |
| 244 | |
| 245 | def get_input(self, batch, k): |
| 246 | x = batch[k] |