Shouldn't predictions accumulate on CPU, not GPU?

[jamarju]

In https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/trainer/predict_loop.py#L111:

        predictions = self.trainer.accelerator.predict_step(args)

        if predictions is None:
            self.warning_cache.warn("predict returned None if it was on purpose, ignore this warning...")

        self.trainer.call_hook("on_predict_batch_end", predictions, batch, batch_idx, dataloader_idx)

        if self.should_store_predictions:
            self.predictions[dataloader_idx].append(predictions)

Shouldn't .cpu() be called on predictions before appending them to self.predictions so that predictions accumulate con CPU memory, not GPU memory?

[tommasobonomo]

This issue has just affected me. I had trained a model using a single GPU, it was tight but no CUDA OOM errors, then I tried getting some predictions through pytorch_lightning.Trainer and because of the accumulation on GPU I got those CUDA OOM.

Would be nice to have a flag (accumulate_on_cpu?) that would accumulate predictions on CPU.

[tchaton]

Dear @jamarju @tommasobonomo,

Any of you would be interested by making a PR to move the predictions to cpu using apply_to_collection ?

Best,
T.C

[EricWiener]

@tchaton I can do it

[jzazo]

has this been implemented or is planned? Thanks!

[rohitgr7]

I believe doing just prediction.cpu() inside predict_step should work just fine. In case it's some sort of other data structure, one can just use predictions = move_data_to_device(predictions, 'cpu') from pl.utilities.

[EricWiener]

This was implemented in #9085

[EricWiener]

This was fixed in #9085