Fix some arguments
@@ -624,8 +624,13 @@ class Chain(Link):
d[name].to_cpu()
return self
- def to_gpu(self, device_id=None):
- with cuda.get_device_from_id(device_id):
+ def to_gpu(self, device=None):
+ if device is None:
+ device_id = cuda.cupy.cuda.get_device_id()
+ device = cuda.get_device_from_id(device_id)
+ elif type(device) in cuda._integer_types:
+ device = cuda.get_device_from_id(device)
+ with device:
super(Chain, self).to_gpu()
d = self.__dict__
for name in self._children:
@@ -778,8 +783,13 @@ class ChainList(Link):
link.to_cpu()
return self
- def to_gpu(self, device_id=None):
- with cuda.get_device_from_id(device_id):
+ def to_gpu(self, device=None):
+ if device is None:
+ device_id = cuda.cupy.cuda.get_device_id()
+ device = cuda.get_device_from_id(device_id)
+ elif type(device) in cuda._integer_types:
+ device = cuda.get_device_from_id(device)
+ with device:
super(ChainList, self).to_gpu()
for link in self._children:
link.to_gpu()
