I have image features’ tensor with 4 dimensions. I want to apply k-means clustering to 3rd and 4th dimension only leaving the first 2 dimensions as is. In short, I want to reduce the the size of 3rd and 4th dimension to 36. At the end I need tensor of size [2048,36]. If someone could guide me how do I do that ?