Is there any pytorch function can combine the specific continuous dimensions of tensor into one?

Question

Let's call the function I'm looking for "magic_combine", which can combine the continuous dimensions of tensor I give to it. For more specific, I want it to do the following thing:

a = torch.zeros(1, 2, 3, 4, 5, 6)  
b = a.magic_combine(2, 5) # combine dimension 2, 3, 4 
print(b.size()) # should be (1, 2, 60, 6)

I know that torch.view() can do the similar thing. But I'm just wondering if there is any more elegant way to achieve the goal?

Answer 1

a = torch.zeros(1, 2, 3, 4, 5, 6)
b = a.view(*a.shape[:2], -1, *a.shape[5:])

Seems to me a bit simpler than the current accepted answer and doesn't go through a list constructor (3 times).