Add device check for training

[Joe-Vincent]

if torch.cuda.is_available(): device = torch.device("cuda")
else: device = torch.device("cpu")

[javieryu]

@Joe-Vincent I did this, and now most of the computation is done on the GPU. Seems to have a minimal speed up vs CPU which makes me think i did something wrong (1m30s on GPU and 1m50s on CPU for the same trial run). Check the new gpu_fix branch for the changes.