RuntimeError: Expected a 'cpu' device type for generator but found 'cuda'

[Six-walnuts-Lu]

hello，How to solve this problem?

Traceback (most recent call last):
  File "main_cola.py", line 207, in <module>
    main()
  File "main_cola.py", line 86, in main
    loader_iter = iter(train_loader)
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 359, in __iter__
    return self._get_iterator()
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 305, in _get_iterator
    return _MultiProcessingDataLoaderIter(self)
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 944, in __init__
    self._reset(loader, first_iter=True)
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 975, in _reset
    self._try_put_index()
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 1209, in _try_put_index
    index = self._next_index()
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 512, in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/sampler.py", line 227, in __iter__
    for idx in self.sampler:
  File "/home/linux01/.local/lib/python3.8/site-packages/torch/utils/data/sampler.py", line 124, in __iter__
    yield from torch.randperm(n, generator=torch.Generator(device='cuda')
RuntimeError: Expected a 'cpu' device type for generator but found 'cuda'