neural ode backpropagation error

[Negar-Erfanian]

I'm running this code for the earthquake data and had to do some readjustments as now I'm also including depth in the dataset so spatial locations are now 3d (instead of being 2d: lat and long)

Now I'm facing this issue which is due to backpropagation through the neural ode (odjoint_ode) and I really have no idea what the problem is. Can someone help me with this?

-- Process 0 terminated with the following error:
Traceback (most recent call last):
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
fn(i, *args)
File "/home/ne12/Downloads/neural_stpp-main/train_stpp.py", line 145, in main
_main(rank, world_size, args, savepath, logger)
File "/home/ne12/Downloads/neural_stpp-main/train_stpp.py", line 370, in _main
loss.backward()
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torch/tensor.py", line 245, in backward
torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torch/autograd/init.py", line 145, in backward
Variable._execution_engine.run_backward(
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torch/autograd/function.py", line 89, in apply
return self._forward_cls.backward(self, *args) # type: ignore
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/adjoint.py", line 126, in backward
aug_state = odeint(
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/odeint.py", line 77, in odeint
solution = solver.integrate(t)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/solvers.py", line 28, in integrate
self._before_integrate(t)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/rk_common.py", line 161, in _before_integrate
f0 = self.func(t[0], self.y0)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/misc.py", line 191, in forward
return self.base_func(t, y)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/misc.py", line 159, in forward
return self.mul * self.base_func(-t, y)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
result = self.forward(*input, **kwargs)
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/misc.py", line 138, in forward
f = self.base_func(t, _flat_to_shape(y, (), self.shapes))
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/adjoint.py", line 93, in augmented_dynamics
_params = tuple(torch.as_strided(param, (), ()) for param in adjoint_params) # noqa
File "/opt2/software/Anaconda3/2020.11/lib/python3.8/site-packages/torchdiffeq/_impl/adjoint.py", line 93, in
_params = tuple(torch.as_strided(param, (), ()) for param in adjoint_params) # noqa
RuntimeError: setStorage: sizes [], strides [], storage offset 0, and itemsize 4 requiring a storage size of 4 are out of bounds for storage of size 0