Hi!
I have a decoder outputing a tensor (let’s call it A) of size [batch_size, num_points, 3], basically a Point-Cloud in three-dimensional euclidean space. I’d like to perform sorting on each of the generated PCs with respect to a specific axis (let’s say z-axis) before feeding the now sorted (named B) [batch_size, num_points, 3] tensor to the next network layers.
The problem arise because I’ve no idea on how to differentiably sort A in one step with torch.sort(). Differently, I tried also to reduce A to a new tensor of size [batch_size, num_points, 1], keeping only z-axis values, performing sort and then using the retrieved indices again on the A tensor. By doing so I’ve noticed that while the values returned by torch.sort have requires_grad=True, the same is not true for indices. How this impact on general differentiability? Is there a better way to sort each of the generated PCs with respect to a single axis (I mean xyz here) without loosing differentiability?
Thank you!