Hi,
I guess you can do
input = input.transpose(dim, -1).contiguous()
return input.view(-1, input.size(-1))

I am not sure how you expect that flattening to happen but this will swap the selected dimension with the last one and flatten all the first dimensions using pytorch’s row major layout.