def test_zerograds_fill_multi_gpu(self):
cupy = cuda.cupy
with cuda.get_device(1):
a = chainer.Variable(cupy.empty(3, dtype=np.float32))
a.grad = cupy.empty_like(a.data)
a.zerograd()
self.assertEqual(int(a.grad.device), 1)
with cuda.get_device(1):
g_expect = cupy.zeros_like(a.data)
cupy.testing.assert_array_equal(a.grad, g_expect)
评论列表
文章目录