'One-hot encodes the waveform amplitudes.
This allows the definition of the network as a categorical distribution
over a finite set of possible amplitudes.'

def _one_hot(self, input_batch):

    with tf.name_scope('one_hot_encode'):
        encoded = tf.one_hot(input_batch, depth=self.quantization_channels, dtype=tf.float32)
        shape = [self.batch_size, (-1), self.quantization_channels]
        encoded = tf.reshape(encoded, shape)
    return encoded
