use generator arg in dataloader in cevae

[martinjankowiak]

i guess this is a torch bug? otherwise you get something like

  File "/home/mjankowiak/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 628, in __next__
    data = self._next_data()
  File "/home/mjankowiak/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 670, in _next_data
    index = self._next_index()  # may raise StopIteration
  File "/home/mjankowiak/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 618, in _next_index
    return next(self._sampler_iter)  # may raise StopIteration
  File "/home/mjankowiak/.local/lib/python3.8/site-packages/torch/utils/data/sampler.py", line 254, in __iter__
    for idx in self.sampler:
  File "/home/mjankowiak/.local/lib/python3.8/site-packages/torch/utils/data/sampler.py", line 132, in __iter__
    yield from torch.randperm(n, generator=generator).tolist()
RuntimeError: Expected a 'cuda' device type for generator but found 'cpu'