NaN Values with deep kernel

[root221]

Hi,
I am attempting to utilize the deep kernel with Bayesian optimization (TurBo). However, I encounter NaN values immediately after initialization, before any training takes place. This issue specifically occurs when using both the RFF kernel and deep kernel, whereas it does not occur when I use the MaternKernel or RBF kernel. Any hints on why this might happen and how to fix it? I have attached my code for reference.

import torch.nn.functional as F
import torch.nn as nn
class FeatureExtractor(torch.nn.Sequential):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 8, kernel_size=3, dtype=torch.float64)
        self.conv2 = nn.Conv2d(8, 8, kernel_size=3, dtype=torch.float64)
        self.fc1 = nn.Linear(128, 2, dtype=torch.float64)
        #self.fc2 = nn.Linear(16, 1, dtype=torch.float64)
                        
    def forward(self, x):
        window_size = int(np.sqrt(x.shape[-1])) 
        x = x.view(-1, 1, window_size, window_size)
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = x.view(-1, 128)
        #x = F.relu(self.fc1(x))
        x = self.fc1(x)
        return x

class GPModelDKL(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        feature_extractor = FeatureExtractor().cuda()
        super(GPModelDKL, self).__init__(train_x, train_y, likelihood)
        
        self.mean_module = gpytorch.means.ConstantMean()
        #self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())
        self.covar_module = gpytorch.kernels.GridInterpolationKernel(
                gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=2)),
                num_dims=2, grid_size=100
        )
        self.likelihood = likelihood
        self.feature_extractor = feature_extractor
        self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)
        
    def forward(self, x):
        x = self.feature_extractor(x)
        x = self.scale_to_bounds(x)
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = GPModelDKL(X_Sobol, Y_Sobol, likelihood).to(X_Sobol.dtype)


if torch.cuda.is_available():
    model = model.cuda()
    likelihood = likelihood.cuda()


model.train()
likelihood.train()

output = model(X_Sobol.cuda())

I would greatly appreciate any advice or hints on how to resolve this issue and successfully run the code without encountering errors.

Thank you in advance for your assistance.