(self, input, return_pred_indices=False)
| 116 | return dec |
| 117 | |
| 118 | def forward(self, input, return_pred_indices=False): |
| 119 | quant, diff, (_,_,ind) = self.encode(input) |
| 120 | dec = self.decode(quant) |
| 121 | if return_pred_indices: |
| 122 | return dec, diff, ind |
| 123 | return dec, diff |
| 124 | |
| 125 | def get_input(self, batch, k): |
| 126 | x = batch[k] |