Error making prediction with pruned neural network

I have trained an MLP for regression with 18 output nodes. I want to remove 17 of these nodes (i.e. the weights and biases) so the MLP can output only for the one remaining node. I successfully removed these nodes, however the new MLP can only predict for one single example at a time (while the original MLP can predict for 10000 examples at a time). When I try to predict more than one example at a time I receive the following error - “RuntimeError: mat2 must be a matrix, got 1-D tensor”

My original network I trained was defined as follows:
class NMR_Model(nn.Module):
** def init(self):**
** super().init()**
** self.lin1 = nn.Linear(14000, 200)**
** self.relu1 = nn.ReLU()**
** self.lin2 = nn.Linear(200, 18)**
** def forward(self, input):**
** return (self.lin2(self.relu1(self.lin1(input))))**

model = NMR_Model()

After training, I made a second model and used the same weights and biases to create the MLP with only one output node. The following is the code for this:
class NMR_Model2(nn.Module):
** def init(self):**
** super().init()**
** self.lin1 = nn.Linear(14000, 200)**
** self.relu1 = nn.ReLU()**
** self.lin2 = nn.Linear(200, 1)**
** def forward(self, input):**
** return (self.lin2(self.relu1(self.lin1(input))))**

model2 = NMR_Model2()
model2.lin1.weight.data = model.lin1.weight.data
model2.lin1.bias.data = model.lin1.bias.data
model2.lin2.weight.data = model.lin2.weight.data[0]
model2.lin2.bias.data = model.lin2.bias.data[0]

The following works to make predictions with the first model, with outputs being the predictions for 10000 examples [i.e. the X_test variable]:
outputs = model(X_test)

This same code does not work for the pruned model, however the following works to make a prediction for a single example:
outputs = model2(X_test[0])

I am wondering why the pruned model only works for a single example, and how I might could modify my code to allow prediction for all 10000 examples using the pruned model.

Figured it out. I saw on this post (RuntimeError: mat2 must be a matrix, got 1-D tensor) that someone had a similar error, and the solution suggested they needed to change the shape of their output tensor. I modified my code where I was copying the weights/biases from the first network, as follows:

model3 = NMR_Model2()
model3.lin1.weight.data = model2.lin1.weight.data
model3.lin1.bias.data = model2.lin1.bias.data
model3.lin2.weight.data = model2.lin2.weight.data[:1]
model3.lin2.bias.data = model2.lin2.bias.data[:1]