Hello,

I'm interested in trying this model, but I'm finding that with standard defaults I get an immediate solution blowup.

For example, if I consider the implementation in health_care/network.py and run the following code

data = torch.randn(10, 50, 3)
model = UnICORNN(3, 10, 1, 0.011, 9.0, 1)
print(model(data))

which are the dt and alpha defaults for the health care task, I get

[ 3.2111e+18],
[-7.2745e+00],
[-2.6744e+23],
[ 4.5934e+32],
[ 2.4417e+28],
[ 4.3936e+27],
[-1.5807e+27],
...

For cupy version 9.x you can just:

import cupy as cp

...

mod = cp.RawModule(code=UnICORNN_CODE, options=('--std=c++11',),
        name_expressions=('unicornn_fwd', 'unicornn_bwd'))
...

which seems to work and you don't need to go through pynvrtc etc.

Which part of the code should I delete and where should I insert mod = cupy.RawModule(code=UnICORNN_CODE, options=('--std=c++11',), name_expressions=('unicornn_fwd', 'unicornn_bwd'))

class UnICORNN_compile():
    _UnICORNN_PROG = Program(UnICORNN_CODE, 'unicornn_prog.cu')
    _UnICORNN_PTX = _UnICORNN_PROG.compile()
    _DEVICE2FUNC = {}

    def __init__(self):
        super(UnICORNN_compile, self).__init__()

    def compile_functions(self):
        device = torch.cuda.current_device()
        mod = function.Module()
        mod.load(bytes(self._UnICORNN_PTX.encode()))
        fwd_func = mod.get_function('unicornn_fwd')
        bwd_func = mod.get_function('unicornn_bwd')

        Stream = namedtuple('Stream', ['ptr'])
        current_stream = Stream(ptr=torch.cuda.current_stream().cuda_stream)

        self._DEVICE2FUNC[device] = (current_stream, fwd_func, bwd_func)
        return current_stream, fwd_func, bwd_func

    def get_functions(self):
        res = self._DEVICE2FUNC.get(torch.cuda.current_device(), None)
        return res if res else self.compile_functions()