Hi how’s it going? I have a simple neural net that takes in a convolutional layer, flattens the output, concats another observation to the output, and then passes it through two linear layers. I know how to do everything in nn.Sequential except for the concatenation. How would I structure this?
Here is the original neural net:
        self.conv1 = nn.Conv2d(in_channels=in_dim, out_channels=1, kernel_size=(1, 1), stride=1)
        self.fc1 = nn.Linear(flattened_dim+obs_dim,20)
        self.fc2 = nn.Linear(20,output_dim)

        def network(x,obs):
            x = self.conv1(x)
            x = torch.flatten(x)
            x = torch.cat((x,obs))
            x = self.fc1(x)
            x = self.fc2(x)

            return x

And I’d like to convert it to something like this:
        self.network = nn.Sequential(
                        nn.Conv2d(in_channels=in_dim, out_channels=1, kernel_size=(1, 1), stride=1),
                        nn.Flatten(start_dim=0),

                       '''Here I need to concatenate an observation with 
                         the output of nn.Flatten before the next layer'''

                        nn.Linear(flattened_dim+obs_dim,20),
                        nn.Linear(20,output_dim),
                    )

How would you go about implementing this?