Fix bug introduced by recent refactoring in multi-GPU implementations