Fix issues #676 and #678

More tests are necessary, specially using a multi-gpu environment. Anyway, everything seems fine.