How to quickly cut slices from cuda tensor wrt to another tensors values

Question

I have a torch cuda tensor A of shape (100, 80000, 4), and another cuda tensor B of shape (1, 80000, 1). What I want to d