Hi,
I’m trying to create a nn.Module using index_select and got stuck in a small problem. The code below does nothing, but shows the issue:
import torch
from torch import nn
from torch.autograd import Variable
import numpy as np


class test(nn.Module):
    def forward(self, x):
        a = torch.LongTensor(np.arange(x.size()[1]))
        return x.index_select(1, a)


if __name__ == '__main__':
    dtype = torch.FloatTensor
    x = Variable(torch.zeros(1, 2, 8, 8).type(dtype), requires_grad=False)

    # Works
    model = torch.nn.Sequential(
        test(),
        nn.Conv2d(2, 2, 3, bias=True),
    )
    print(model)
    print(model(x))

    # Works
    model = torch.nn.Sequential(
        test(),
        test(),
        nn.Conv2d(2, 2, 3, bias=True),
    )
    print(model)
    print(model(x))

    # Doesn't work
    model = torch.nn.Sequential(
        nn.Conv2d(2, 2, 3, bias=True),
        test(),
        nn.Conv2d(2, 2, 3, bias=True),
    )
    print(model)
    print(model(x))

As shown in the code, if I use the new module combined with Conv2d a problem happens (only) when I put a convolutional layer before the “test” new module (commented “Doesn’t work”).  I get the following error message:
return IndexSelect.apply(self, dim, index)

RuntimeError: save_for_backward can only save input or output tensors, but argument 0 doesn't satisfy this condition

I would appreciate if you could help me out.
Thanks for the great software!