Refactored a bit to introduce multiple gpu handling