neural networks in pytorch
I am using pytorch and need to implement this as a part of a neural network. Is there a particular way to code the layers shown in purple (s
def forward(self, images: torch.Tensor) -&gt; torch.Tensor:
x= self.fc1(x)
return x

Combining all my comments into the answer.
To split the output vector from the first layer, which has a shape [batch_size, 4608], you can use torch.split function as follows
batch = torch.zeros((10, 4608))
sub_batch_1, sub_batch_2 = torch.split(batch, 2304, dim=1)
print(sub_batch_1.shape, sub_batch_2.shape)

This code results in two tensors
(torch.Size([10, 2304]), torch.Size([10, 2304]))

I am not sure about MAX logic you need. If it is about getting the maximum element in a given position from two tensors obtained during split, then torch.maximum function can be used to do so.
