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.