We can directly test if backprop works with both slice and gather:
import torch

class SliceNet(torch.nn.Module):
    def __init__(self):
        super(SliceNet, self).__init__()
        self.linear1 = torch.nn.Linear(32, 64)
        self.sigmoid = torch.nn.Sigmoid()
        self.linear2 = torch.nn.Linear(32, 1, bias=True)

    def forward(self, x):
        x = self.linear1(x)
        x = self.sigmoid(x)
        x = x[:,:32]
        x = self.linear2(x)
        return self.sigmoid(x)

class GatherNet(torch.nn.Module):
    def __init__(self):
        super(GatherNet, self).__init__()
        self.linear1 = torch.nn.Linear(32, 64)
        self.sigmoid = torch.nn.Sigmoid()
        self.linear2 = torch.nn.Linear(32, 1, bias=True)
        batch_size = 2
        self.indices = torch.arange(0, 64, 2).repeat(batch_size, 1)

    def forward(self, x):
        x = self.linear1(x)
        x = self.sigmoid(x)
        x = torch.gather(x, 1, self.indices)
        x = self.linear2(x)
        return self.sigmoid(x)

def test_gather_gradient():
    print("testing gathernet, loss should go down")
    data = torch.randn(2, 32)
    label = torch.randn(2, 1).round()
    gathernet = GatherNet()
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(gathernet.parameters(), 1e-4)
    for i in range(10000):
        out = gathernet(data)
        loss = criterion(out, label)
        if i % 1000 == 0:
            print(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

def test_slice_gradient():
    print("testing slicenet, loss should go down")
    data = torch.randn(2, 32)
    label = torch.randn(2, 1).round()
    slicenet = SliceNet()
    criterion = torch.nn.MSELoss()
    optimizer = torch.optim.SGD(slicenet.parameters(), 1e-4)
    for i in range(10000):
        out = slicenet(data)
        loss = criterion(out, label)
        if i % 1000 == 0:
            print(loss.item())
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

test_gather_gradient()
test_slice_gradient()

testing gathernet, loss should go down
0.16012126207351685
0.13146084547042847
0.10953329503536224
0.09261573851108551
0.07940082252025604
0.06893321871757507
0.06052353233098984
0.053674302995204926
0.04802417382597923
0.043307825922966
testing slicenet, loss should go down
1.1417878866195679
0.950451672077179
0.828807532787323
0.7515791654586792
0.7005079984664917
0.6650912761688232
0.6394497156143188
0.6201967597007751
0.60529625415802
0.593470573425293

In general, I don’t think most indexing operations should be problematic for backprop. For example, could max pooling work if indexing was a problem?