__init__() got an unexpected keyword argument 'func'

[ClaudiaShu]

Hi, I am trying to implement neuralSDE using torchdyn, my code is:

# drift_func
f = nn.Sequential(...)
# diffusion function
g = nn.Sequential(...)                           
self.func = NeuralSDE(f, g, solver=args.solver, rtol=args.rtol, atol=args.atol)

But I got this error:

TypeError: __init__() got an unexpected keyword argument 'func'

I checked the source code and it seems that the initialisation of the model class is different from what input it takes.

class NeuralSDE(SDEProblem, pl.LightningModule):
    def __init__(self, drift_func, diffusion_func, noise_type ='diagonal', sde_type = 'ito', order=1,
                 sensitivity='autograd', s_span=torch.linspace(0, 1, 2), solver='srk',
                 atol=1e-4, rtol=1e-4, ds = 1e-3, intloss=None):
        super().__init__(func=SDEFunc(f=drift_func, g=diffusion_func, order=order), order=order, sensitivity=sensitivity, s_span=s_span, solver=solver,
                                      atol=atol, rtol=rtol)

What should be given to the model as the input or is there a bug that remains to fix?

Thanks in advance!