DDP Trainer's test method -> TypeError: can't pickle SwigPyObject objects

[sshleifer]

I call
my code (roughly)

module = pl.Module(...)
trainer = pl.Trainer(module, distributed_backend='ddp', n_gpu=2,...)
trainer.fit() # works fine uses all GPUs
trainer.test(model) # code works only with n_gpu=1 or n_gpu=0.

Traceback:

    trainer.test(model)
../miniconda3/envs/nb/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py:1064: in test
    self.fit(model)
../miniconda3/envs/nb/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py:844: in fit
    mp.spawn(self.ddp_train, nprocs=self.num_processes, args=(model,))
../miniconda3/envs/nb/lib/python3.7/site-packages/torch/multiprocessing/spawn.py:200: in spawn
    return start_processes(fn, args, nprocs, join, daemon, start_method='spawn')
../miniconda3/envs/nb/lib/python3.7/site-packages/torch/multiprocessing/spawn.py:149: in start_processes
    process.start()
../miniconda3/envs/nb/lib/python3.7/multiprocessing/process.py:112: in start
    self._popen = self._Popen(self)
../miniconda3/envs/nb/lib/python3.7/multiprocessing/context.py:284: in _Popen
    return Popen(process_obj)
../miniconda3/envs/nb/lib/python3.7/multiprocessing/popen_spawn_posix.py:32: in __init__
    super().__init__(process_obj)
../miniconda3/envs/nb/lib/python3.7/multiprocessing/popen_fork.py:20: in __init__
    self._launch(process_obj)
../miniconda3/envs/nb/lib/python3.7/multiprocessing/popen_spawn_posix.py:47: in _launch
    reduction.dump(process_obj, fp)
../miniconda3/envs/nb/lib/python3.7/multiprocessing/reduction.py:60: in dump
    ForkingPickler(file, protocol).dump(obj)
E   TypeError: can't pickle SwigPyObject objects

Is there something I can do to make this work?
Thanks!

[sshleifer]

This is related, I think: https://discuss.pytorch.org/t/how-to-debug-saving-model-typeerror-cant-pickle-swigpyobject-objects/66304/3