def set_state(self, state): self.x, self.c, self.h = F.split_axis( state, (self.in_size, self.in_size + self.size), axis=1)