If I understand your use case correctly, this should work:
B, N, D = 2, 3, 4
hidden = torch.randn(B, N, D)
y = torch.zeros(B, D)
eot_inds = torch.randint(0, N, (B,))

for i in range(B):
    y[i] = hidden[i, eot_inds[i]]
    

out = hidden[torch.arange(hidden.size(0)), eot_inds]
print((out == y).all())
> tensor(True)