How to chunking a 2-D tensor

3

Thanks !
This is what I need !

I want to ask a related question.
What if I need to torch.sum() for each chunk, is there any way to avoid the iteration ?