How to add new inputs and features in model FNO 2d

Hi everyone!
Guys, I am new in pytorch and i try to adopt model FNO2d(ref: neuraloperator/fourier_2d_time.py at master · neuraloperator/neuraloperator · GitHub)

Task:
I’m trying to predict the pressure field in oil production.
I have two wells for extraction and injection. In its primary form, the model FNO includes only a field without properties(additional inputs) and I’m trying to add these properties like: porosity, permeability and well location.

First training data set had size [Case = 300,dim1 = 64, dim2 = 86, timesteps= 84], after adding new properties into data size changed to [300, 64, 86, 84, 4] (under 0 - pressure feld).

And this is how I did it:

        self.modes1 = modes1
        self.modes2 = modes2
        self.width = width
        self.padding = 8 # pad the domain if input is non-periodic

        self.p = nn.Linear(15, self.width) # input channel is 15: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y) + 3 properties
        self.conv0 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.conv1 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.conv2 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.conv3 = SpectralConv2d(self.width, self.width, self.modes1, self.modes2)
        self.mlp0 = MLP(self.width, self.width, self.width)
        self.mlp1 = MLP(self.width, self.width, self.width)
        self.mlp2 = MLP(self.width, self.width, self.width)
        self.mlp3 = MLP(self.width, self.width, self.width)
        self.w0 = nn.Conv2d(self.width, self.width, 1)
        self.w1 = nn.Conv2d(self.width, self.width, 1)
        self.w2 = nn.Conv2d(self.width, self.width, 1)
        self.w3 = nn.Conv2d(self.width, self.width, 1)
        self.norm = nn.InstanceNorm2d(self.width)
        self.q = MLP(self.width, 1, self.width * 4) # output channel is 1: u(x, y)

    def forward(self, x):

        grid = self.get_grid(x.shape, x.device)
        pressure = self.get_Pressure(x, x.shape, x.device)
        properties = self.get_Properties(x, x.shape, x.device)

        x = torch.cat((pressure, grid, properties), dim=-1)
        x = self.p(x)
        x = x.permute(0, 3, 1, 2)
#....and so on

But it doesn’t work correctly. When i try predict pressure with new well location, it has too big loss and visual also not correct. But when wells location same like in training data(in the training dataset wells had fixed location for all cases) it work correctly with some noises.
Example:

I would be grateful for any help or advice.