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.