speed improvement

bug correction.
Should improve the speed on compute 3.5 as the gpu usage is now close to
100%