def init_state(self, param):
xp = cuda.get_array_module(param.data)
with cuda.get_device_from_array(param.data):
self.state['ms'] = xp.zeros_like(param.data)
微信
公众号
扫码关注公众号