Fix issues #676 and #678

More test is necessary, specially using a multi-gpu environment.