Passing device to GNNModel constructor

[luisaforozco]

For the moment we leave like this, but If we manage to pass the scaling factors as vectors that are in device from train_one_step_wandb.py, then we can remove the parameter in the constructor.

Originally posted by @luisaforozco in #3 (comment)