How can I train models with differnt GPU?

In the old version, I can use torch.cuda.set_device(index). But in 1.0 it might not work.

epoch train loss valid loss accuracy
THCudaCheck FAIL file=/opt/conda/conda-bld/pytorch-nightly_1539421063549/work/aten/src/THC/generic/THCTensorCopy.cpp line=21 error=77 : an illegal memory access was encountered
Traceback (most recent call last):
File “runsub.py”, line 27, in
learn.fit_one_cycle(10,max_lr=0.01)
File “/home/user/miniconda2/envs/fastai/lib/python3.7/site-packages/fastai/train.py”, line 18, in fit_one_cycle
learn.fit(cyc_len, max_lr, wd=wd, callbacks=cbs)
File “/home/user/miniconda2/envs/fastai/lib/python3.7/site-packages/fastai/basic_train.py”, line 133, in fit
callbacks=self.callbacks+callbacks)
File “/home/user/miniconda2/envs/fastai/lib/python3.7/site-packages/fastai/basic_train.py”, line 86, in fit
raise e
File “/home/user/miniconda2/envs/fastai/lib/python3.7/site-packages/fastai/basic_train.py”, line 71, in fit
xb, yb = cb_handler.on_batch_begin(xb, yb)
File “/home/user/miniconda2/envs/fastai/lib/python3.7/site-packages/fastai/callback.py”, line 198, in on_batch_begin
a = cb.on_batch_begin(**self.state_dict)
File “/home/user/miniconda2/envs/fastai/lib/python3.7/site-packages/fastai/callbacks/mixup.py”, line 17, in on_batch_begin
lambd = last_input.new(lambd)
RuntimeError: cuda runtime error (77) : an illegal memory access was encountered at /opt/conda/conda-bld/pytorch-nightly_1539421063549/work/aten/src/THC/generic/THCTensorCopy.cpp:21

Just change the value of defaults.device to the gpu you want.