Can RevTorch work in multi-gpu?

RobinBruegger / RevTorch

Framework for creating (partially) reversible neural networks with PyTorch

BSD 3-Clause "New" or "Revised" License

145 stars 18 forks source link

Can RevTorch work in multi-gpu? #10

Open ewrfcas opened 4 years ago

ewrfcas commented 4 years ago

When I use revtorch with DataParallel, the GPU memory is overflow after the backward.

RobinBruegger commented 4 years ago

Hi @ewrfcas

Unfortunately I do not have access to multiple GPUs to test/debug this behaviour. I would however be happy to fix it if you could provide a pull request for this issue