You could try it with this code sample:
A = torch.randn(10, 10)
mask = torch.ByteTensor(10, 10).random_(2)
b = torch.masked_select(A, mask)

C = torch.zeros(10, 10)
C[mask] = b
print(A, mask, C)

I’m not sure if it fits your needs regarding the functionality with auto-grad.
Could you test it or give a small code sample showing what you would like to do?