def check_accumulate_grads_from_gpu(self, src_id):
with cuda.Device(src_id):
self.optimizer.accumulate_grads([cuda.cupy.arange(3)])
grad = self.target.param.grad
self.assertTrue((cuda.to_cpu(grad) == np.arange(3) * 2).all())
评论列表
文章目录