I’m trying to build a neural network includes an SDEs layer and full connected NN. The problem that it takes hours just to complete one epoch.

This is my code:
First the SDEs class

class SDE(nn.Module):

def __init__(self):
    super().__init__()
    self.theta = nn.Parameter(torch.tensor(0.1), requires_grad=True)  # Scalar parameter.
    self.noise_type = "diagonal"
    self.sde_type = "ito"

def f(self, t, y):
    return torch.sin(t) + self.theta * y

def g(self, t, y):
    return 0.3 * torch.sigmoid(torch.cos(t) * torch.exp(-y))

this is the pyTorch model:

class SDENET__(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
    super(SDENET_, self).__init__()
    self.batch_size, self.state_size, self.t_size = 1, 1, 16
    self.ys = torch.tensor(0.1)
    self.fc = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
    sde = SDE()
    ts = torch.linspace(0, 1, self.t_size)
    y0 = torch.full(size=(self.batch_size, self.state_size), fill_value=0.1)
    with torch.no_grad():
        self.ys = torchsde.sdeint(sde, y0, ts, method='euler')
    ys = self.ys[:, -1, :]
    out = self.fc(ys)
    return out

the model is to train NN for multivariate timeseries forecating.