Custom gather method

[Rizhiy]

How do I correctly override gather behaviour?
Ideally, I would have a single place where I should specify new method and it would work with dp, ddp and ddp_spawn.

[Rizhiy]

I'm currently using accelerator="dp" and getting the following error:

...
    return trainer.predict(
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:631: in predict
    results = self._run(model)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:756: in _run
    self.dispatch()
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:795: in dispatch
    self.accelerator.start_predicting(self)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py:102: in start_predicting
    self.training_type_plugin.start_predicting(trainer)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py:152: in start_predicting
    self._results = trainer.run_stage()
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:806: in run_stage
    return self.run_predict()
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py:1083: in run_predict
    self.predict_loop.predict_step(batch, batch_idx, dataloader_idx)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/trainer/predict_loop.py:111: in predict_step
    predictions = self.trainer.accelerator.predict_step(args)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py:265: in predict_step
    return self.training_type_plugin.predict_step(*args)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/dp.py:107: in predict_step
    return self.model(*args, **kwargs)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/torch/nn/modules/module.py:889: in _call_impl
    result = self.forward(*input, **kwargs)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py:168: in forward
    return self.gather(outputs, self.output_device)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/torch/nn/parallel/data_parallel.py:180: in gather
    return gather(outputs, output_device, dim=self.dim)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py:76: in gather
    res = gather_map(outputs)
../../miniconda3/envs/ntf/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py:71: in gather_map
    return type(out)(map(gather_map, zip(*outputs)))
TypeError: 'MyModel' object is not iterable

I need trainer to use my custom gather_map instead of the one in torch.

[Rizhiy]

Managed to get this working for dp by subclassing DataParallelPlugin and passing it to trainer. Although this approach requires individual plugins for each accelerator, it would be nice to have a way to set this for all accelerators at the same time.

[ananthsub]

@Rizhiy mind commenting on #7534 for what your desired behavior would look like?