This might work:
block, example, label = 2, 3, 4
N = 5

data = torch.randn(block, example, label)
index = torch.empty(N, block, dtype=torch.long).random_(label)
res = torch.cat([torch.gather(data, 2, i[:, None, None].expand(-1, example, -1)) for i in index])
res = res.view(N, block, example)

Could you check this code with your data?
I’m not sure about the last .view operation and also if it’s possible to get rid of the torch.cat and the list comprehension.