delete JIT distance_module before replicating
@@ -59,6 +59,13 @@ class MultiDeviceKernel(DataParallel, Kernel):
if len(self.device_ids) == 1:
return self.module.forward(*inputs[0], **self._kwargs[0])
+ # JIT modules can't be pickled and replicated yet
+ # But reinitializing the distance_module every forward pass
+ # is slow and should be removed once JIT modules can be pickled
+ def set_distance_module_to_none(module):
+ if hasattr(module, 'distance_module'):
+ module.distance_module = None
+ self.module.apply(set_distance_module_to_none)
# 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)])
