export CUDA_VISIBLE_DEVICES=0 before running the code.