compress two consecutive dimension in PyTorch
I'm looking for such a function Tensor.compress(*dims) where dims=(int ...) is a consecutive sequence of integer:
a = torch.rand(2,2,3)
b = a.compress(0,1)

b.size()
&gt;&gt;&gt; (4,3)

I know view would work, however, in case I don't know the shape of a in advance, I have to do an extra operation to acquire its size and then do view, which is not what I want.

You do not have to explicitly &quot;do the math&quot;, torch.view can do some of it for you if you use -1 as the shape of one of the dimensions:
b = a.view(-1, *a.shape[2:])
b.shape
&gt;&gt;&gt; torch.Size([4, 3])

