Fix case where len(device_ids)==1
@@ -43,10 +43,10 @@ class MultiDeviceKernel(DataParallel, Kernel):
inputs = tuple((x1_[0], x2_) for x1_, x2_ in zip(self._x1_scattered, self._x2_subs))
if not self.device_ids:
- return self.module(*inputs, **self._kwargs)
+ return self.module.forward(*inputs, **self._kwargs)
if len(self.device_ids) == 1:
- return self.module(*inputs[0], **self._kwargs[0])
+ return self.module.forward(*inputs[0], **self._kwargs[0])
# Can't cache the replication because the base kernel module can change every time (e.g. param updates)
replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
