PyTorch tensor to List of Lists

I have a PyTorch tensor of shape ((1,512,16,3)).
The (16,3) part of the tensor represents 16 - (x,y,z) points.
These 16 points could be unique points or repeated.
I want to convert it to a list of 512 lists only with unique points.
Dimensions of the resulting list of lists (1,512,p,3) ==> 1 list which has 512 lists each of varying “p” size which in turn contain a list of 3 points.

So far I tried using two for loops to fetch the (16,3) points and use torch.unique to pick the unique points. But 2 for loops make my code slower.

Please let me know if there is a function to do it in PyTorch.

Thanks in advance.

If I understand the use case correctly, p won’t be a fixed number, but is varying for each of the 512 entries, e.g.:

x = torch.randn(1, 512, 16, 3)
# your operation
result = ...
print(result[0, 0].shape) # [4, 3]
print(result[0, 1].shape) # [7, 3]
...

If that’s the case, you won’t be able to create the result as a tensor without padding.

Let me know, if I misunderstood the use case.

You understood it correct.
Thanks a lot. I am using padding now.