Torch SDE simulation error

I tried simulating the Ornstein Uhlenbeck process from torch sde. The SDE model looks like:
$ dX_t = \theta (\mu- X_t)dt + \sigma dW_t $
Here is my code: class OU(nn.Module):

noise_type = "diagonal"
sde_type = "ito"

def __init__(self, theta, mu, sigma):
    self.theta = theta = mu
    self.sigma = sigma
def f(self, t, y):
    return self.theta*( - y)

def g(self, t, y):
    return self.sigma``

When an object is created from this class, I get the following error:

AttributeError                            Traceback (most recent call last)
<ipython-input-25-775824621def> in <module>()
      8 with torch.no_grad():
----> 9     ys = torchsde.sdeint(sde, y0, ts, method='milstein')
     10 plot(ts, ys, xlabel='$t', ylabel='$Y_t$')

1 frames
/usr/local/lib/python3.7/dist-packages/torchsde/_core/ in check_contract(sde, y0, ts, bm, method, adaptive, options, names, logqp)
    205     if hasattr(sde, 'g'):
    206         has_g = True
--> 207         g_diffusion_shape = tuple(sde.g(ts[0], y0).size())
    208         _check_2d_or_3d('Diffusion', g_diffusion_shape)
    209     if hasattr(sde, 'f_and_g'):

**AttributeError: 'float' object has no attribute 'size'**

self.sigma seems to be a float while a tensor is expected as .size() is called on this return value.
I don’t know if transforming the float scalar into a tensor would work or if more changes are needed, but it might be worth a try.

Thanks for the tip! Transforming the self.sigma scalar to a tensor by self.sigma*torch.ones(y.size()) worked for me :smile: