Make sure code works on cpu 
