Hi @Hdk
I has a similar use case where a tensor lets say A, A.shape = (100,10,5,5)
I wanted to add the last two dimensions such that new Ashape = (100,10,5) and following did the trick

A.sum(axis=2)

Hopefully it give you more ideas.
regards,