Gpytorch with error "Shapes are not broadcastable for mul operation"

Question

I am trying to follow the exact deep kernel learning example provided in the GPyTorch documentation to develop my own code. Here is my code:

class CNN_embedding(nn.Module):
  def __init__(self, inplace=False):
    super(CNN_embedding, self).__init__()
    self.conv1 = nn.Conv1d(in_channels = 1, out_channels = 32, kernel_size = 3, stride = 1, padding=0)
    self.relu1 = nn.ReLU()
    self.pool1 = nn.MaxPool1d(kernel_size=2,stride=1)
    self.conv2 = nn.Conv1d(in_channels = 32, out_channels = 16, kernel_size = 3, stride = 1, padding=0)
    self.relu2 = nn.ReLU()
    self.pool2 = nn.MaxPool1d(kernel_size=2,stride=1)
    self.conv3 = nn.Conv1d(in_channels = 16, out_channels = 8, kernel_size = 3, stride = 1, padding=0)
    self.relu3 = nn.ReLU()
    self.pool3 = nn.MaxPool1d(kernel_size=2,stride=1)
    self.flatten = nn.Flatten()
    # self.fc1 = nn.Linear(16, 1)
  def forward(self, x):
    layer1 = self.pool1(self.relu1(self.conv1(x)))
    layer2 = self.pool2(self.relu2(self.conv2(layer1)))
    layer3 = self.pool3(self.relu3(self.conv3(layer2)))
    out = self.flatten(layer3)
    return out

feature_embedding = CNN_embedding()

# We will use the simplest form of GP model, exact inference
class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, X_train, y_train, likelihood):
        super(ExactGPModel, self).__init__(X_train, y_train, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
        self.feature_extractor = feature_embedding
        self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)

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

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = ExactGPModel(X_train, y_train, likelihood)

import os
training_iter = 100


# Find optimal model hyperparameters
model.train()
likelihood.train()

# Use the adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)  # Includes GaussianLikelihood parameters

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

for i in range(training_iter):
    # Zero gradients from previous iteration
    optimizer.zero_grad()
    # Output from model
    output = model(X_train)
    # Calc loss and backprop gradients
    loss = -mll(output, y_train)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f ' % (
        i + 1, training_iter, loss.item() ))
    optimizer.step()


model.eval()
likelihood.eval()

preds = model(X_test)

The input, X_train and X_test, of the model will have shape of (1344,1,10) and (672,1,10), and y_train and y_test will have shape of (1344) and (672). There is no problem to train the model, but when goes to evaluate the model, the error code shows:

Traceback (most recent call last):

  File E:\ACSE\Project\UniformBound\CNN_GP\untitled0.py:159 in <module>
    preds = model(X_test)

  File E:\Python\envs\gpytorch\lib\site-packages\gpytorch\models\exact_gp.py:299 in __call__
    batch_shape = _mul_broadcast_shape(batch_shape, input.shape[:-2])

  File E:\Python\envs\gpytorch\lib\site-packages\gpytorch\utils\broadcasting.py:20 in _mul_broadcast_shape
    raise RuntimeError("Shapes are not broadcastable for mul operation")

RuntimeError: Shapes are not broadcastable for mul operation

It appears there may be an issue with the dimensions of the inputs or with the y_train and y_test. I would greatly appreciate any suggestions or guidance anyone might provide.