Wrote training g CUDA kernel wrapper function. Need to write kernel