Solved by ptrblck in post #2 
              
              
                You could write a custom module and pass add it to nn.Sequential. 
Something like this might work: 
class Narrow(nn.Module):
    def __init__(self, dim, start, length):
        super(Narrow, self).__init__()
        self.dim = dim
        self.start = start
        self.length = length
    
    def …