Hi,

I am trying to train a simple NN which takes 2-d tensor as input data, and outputs a 2-d tensor. Specifically, I would like to have an input and output of shape 16x2.

Below is the model I used.

```
class MyModel(nn.Module):
def __init__(self):
super(MyModel, self).__init__()
self.fc1 = nn.Linear(16, 256)
self.fc2 = nn.Linear(256, 512)
self.fc3 = nn.Linear(512, 256)
self.fc4 = nn.Linear(256, 16)
def forward(self, x):
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = F.relu(self.fc3(x))
x = F.sigmoid(self.fc4(x))
return x
```

With a batch size of 1, I get this error message:

```
RuntimeError: mat1 and mat2 shapes cannot be multiplied (16x2 and 16x256)
```

And, with a batch size of 128 (my preferred option is this, as I need to train my model), I get this error message:

```
RuntimeError: mat1 and mat2 shapes cannot be multiplied (2048x2 and 16x256)
```

Would you please help me how to use a batch of tensors with dimensions 16x2 as input and having a batch of tensors with dimensions 16x2 as output.

Obviously, if batch size is 128, input and output dimensions should be 128x16x2.

Best wishesâ€¦