Mutlitple arguments into ExactGP.forward

[skiyohara]

Dear gpytorch developers and users,

I'm writing a code where a deep neural network and Gaussian process are combined.
I am using as a reference of the code in https://docs.gpytorch.ai/en/latest/examples/06_PyTorch_NN_Integration_DKL/KISSGP_Deep_Kernel_Regression_CUDA.html, but its ExactGP.forward receive an argument. In my case, I want ExactGP.forward to receive multiple arguments.
Here is the pseudocode.

class GPRegressionModel(gpytorch.models.ExactGP):
        def __init__(self, train_x, train_y, likelihood):
            super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
            self.mean_module = gpytorch.means.ConstantMean()
            self.covar_module = gpytorch.kernels.GridInterpolationKernel(
                gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=num_dims)),
                num_dims=num_dims
            )
            self.feature_extractor = feature_extractorNN() # feature_extractorNN is a deep neural network which transform multiple arguments into high dimensions (defined elsewhere)

            self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)

        def forward(self, x,y,z):  # receive multiple arguments
            # We're first putting our data through a deep net (feature extractor)
            projected_x = self.feature_extractorNN(x,y,z)
            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)

During training, namely model.train(), I didn't get any error.
However, when the model was in evaluation mode, namely model.eval(), I got an error below.

model.eval()
pred = model(x=test_x, y=test_y, z=test_z)

Traceback (most recent call last):
File "C:\Users\ku\PycharmProjects\bo_crest\exec\exec_bonn.py", line 281, in
output = model(x = x_test)
File "C:\Users\ku\PycharmProjects\bo_crest\venv\lib\site-packages\gpytorch\models\exact_gp.py", line 284, in call
train_output = super().call(*train_inputs, **kwargs)
File "C:\Users\ku\PycharmProjects\bo_crest\venv\lib\site-packages\gpytorch\module.py", line 30, in call
outputs = self.forward(*inputs, **kwargs)
TypeError: forward() got multiple values for argument 'x'

I think ExactGP.foward implicitly receives "train_x" as an additional input when model.eval(), resulting in that argument 'x' got two arguments, train_x and test_x.
It should be the cause of the error, but I cannot fix it.

Any information would be helpful.
Thanks in advance.

Best regards,
Shin

[skiyohara]

I fixed it.