fixed gpu transfer issue in grad method
fixed grad method
made a test less sensible