GPytorch: RuntimeError: Flattening the training labels failed

Question

I am trying to use GPR with GPytorch, and my function seems to fail when I pass the testing data through the model, but works fine with the training data. The training tensor is of shape (116,6) and the testing tensor is (51,6). The exact error I get is: "RuntimeError: Flattening the training labels failed. The most common cause of this error is that the shapes of the prior mean and the training labels are mismatched. The shape of the train targets is torch.Size([116, 1]), while the reported shape of the mean is torch.Size([116])."

My conversion to tensors is as below:

# convert X_train and X_test to torch arrays

X_train_tensor = torch.tensor(X_train.values, dtype=torch.float64)
X_test_tensor = torch.tensor(X_test.values, dtype=torch.float64)
y_train_tensor = torch.tensor(y_train.values, dtype=torch.float32).reshape(-1, 1)
y_test_tensor = torch.tensor(y_test.values, dtype=torch.float32).reshape(-1, 1)

the class I made originally is this:

class ExactGPModel(gpytorch.models.ExactGP):
#     # Initialize the parameter
    def __init__(self, X_train_tensor, y_train_tensor, likelihood):
        super(ExactGPModel, self).__init__(X_train_tensor, y_train_tensor, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

ChatGPT suggested changing the last line of the class to this but this did not work:

        # return gpytorch.distributions.MultivariateNormal(mean_x.view(-1), covar_x)