- the chunk_wise kernel is used for training - perhaps needs to be changes for the inference of the model