6.0.0b2 does not support to_cpu/to_gpu of inner links

I found that chainer==6.0.0b2breaks my code.

I’m using a helper class that contains a numpy/cupy array and do something with it. The device of this array can be switched by to_cpu/to_gpu method. The link class that uses this helper class overrides to_cpu/to_gpu method to call the helper’s ones.

In chainer==6.0.0b2, a parent link does not call to_cpu/to_gpu of child links. So the following code does not work.

import chainer
from chainer.backends import cuda
import numpy as np


class Helper(object):

    def __init__(self):
        self._value = np.array((0, 1, 2))

    def meth(self, h2):
        # caluculate y from h2 and self._value.
        return y

    def to_cpu(self):
        self._value = cuda.to_cpu(self._value)

    def to_gpu(self, device=None):
        self._value = cuda.to_gpu(self._value, device=device)


class SubModel(chainer.Chain):

    def __init__(self):
        super().__init__()
        with self.init_scope():
            # add some child links

        self._helper = Helper()

    def __call__(self, h1):
        # caluculate h2 from h1 using own child links
        return h2

    def to_cpu(self):
        super().to_cpu()
        self._helper.to_cpu()

    def to_gpu(self, device=None):
        super().to_gpu(device)
        self._helper.to_gpu(device=device)

    def predict(self, h1):
        h2 = self(h1)
        y = self._helper.meth(h2.array)
        return y


class Model(chainer.Chain):

    def __init__(self):
        super().__init__()
        with self.init_scope():
            self.sub_model = SubModel()
            # add some child links

    def __call__(self, x):
        # caluculate h1 from x using own child links
        h2 = self.sub_model(h1)
        return h2

    def predict(self, x):
        # caluculate h1 from x using own child links
        y = self.sub_model.predict(h1)
        return y

[Hakuyume]

I confirmed this issue was fixed in the master branch. Thank you!

[toslunar]

The code in #6825 is a simplification (by @kmaehashi and me) of the code here to reproduce a specific error. Let’s keep this issue open until @Hakuyume confirms it.