title: Can pytorch's autograd handle torch.cat?
body: <p>I'm trying to implement a simple neural network that is supposed to learn an grayscale image. The input consist of the 2d indices of a pixel, the output should be the value of that pixel.</p>
  
  <p>The net is constructed as follows: Each neuron is connected to the input (i.e. the indices of the pixel) as well as to the output of each of the previous neurons. The output is just the output of the last neuron in this sequence.</p>
  
  <p>This kind of network has been very successfull in learning images, as demonstrated e.g. <a href="https://codegolf.stackexchange.com/a/149835/24877">here</a>.</p>
  
  <p><strong>The Problem:</strong> 
  In my implementation the loss function stays between <code>0.2</code> and <code>0.4</code> depending on the number of neurons, the learning rate and the number of iterations used, which is pretty bad. Also if you compare the output to what what we've trained it for there it just looks like noise. But this is the first time I used <code>torch.cat</code> within the network, so I'm not sure whether this is the culprit. Can anyone see what I'm doing wrong?</p>
  
  <pre><code>from typing import List
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  import torch.optim as optim
  from torch.nn import Linear
  
  class My_Net(nn.Module):
      lin: List[Linear]
  
      def __init__(self):
          super(My_Net, self).__init__()
          self.num_neurons = 10
          self.lin = nn.ModuleList([nn.Linear(k+2, 1) for k in range(self.num_neurons)])
  
      def forward(self, x):
          v = x
          recent = torch.Tensor(0)
          for k in range(self.num_neurons):
              recent = F.relu(self.lin[k](v))
              v = torch.cat([v, recent], dim=1)
          return recent
  
      def num_flat_features(self, x):
          size = x.size()[1:]
          num = 1
          for i in size():
              num *= i
          return num
  
  my_net = My_Net()
  print(my_net)
  
  #define a small 3x3 image that the net is supposed to learn
  my_image = [[1.0, 1.0, 1.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]] #represents a T-shape
  my_image_flat = []    #output of the net is the value of a pixel
  my_image_indices = [] #input to the net is are the 2d indices of a pixel
  for i in range(len(my_image)):
      for j in range(len(my_image[i])):
          my_image_flat.append(my_image[i][j])
          my_image_indices.append([i, j])
  
  #optimization loop
  for i in range(100):
      inp = torch.Tensor(my_image_indices)
  
      out = my_net(inp)
  
      target = torch.Tensor(my_image_flat)
      criterion = nn.MSELoss()
      loss = criterion(out.view(-1), target)
      print(loss)
  
      my_net.zero_grad()
      loss.backward()
      optimizer = optim.SGD(my_net.parameters(), lr=0.001)
      optimizer.step()
  
  print("output of current image")
  print([[my_net(torch.Tensor([[i,j]])).item() for i in range(3)] for j in range(3)])
  print("output of original image")
  print(my_image)
  </code></pre>