test_optimizers_by_linear_model.py 文件源码

def test_model_setup_multi_gpu(self):
        with cuda.Device(0):
            model = self.model.model
            optimizer = self.model.optimizer
            model.to_gpu(1)
            optimizer.setup(model)
        for name, param in optimizer.target.namedparams():
            for v in six.itervalues(optimizer._states[name]):
                self.assertEqual(int(param.data.device), int(v.device))