My dataset is small, and I want to load all my dataset into GPU memory when a dataset is created. Meanwhile, I still want to use torch.utils.data.DataLoader because of compatibility with other situations where I load my data on the fly.
My short working example is as follows.
import numpy as np
from torch.utils.data import TensorDataset, DataLoader
import torch

data = np.array([[1,2,3], [4,5,6]])
# I move dataset to GPU first
ds = TensorDataset(torch.Tensor(data).cuda())                                 
dl = DataLoader(ds, batch_size=1, num_workers=1, shuffle=True) 
for x in dl:
    print(x)

However, it crashes.
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-10-11b3bb8f6574> in <module>
      6 ds = TensorDataset(torch.Tensor(data).cuda())
      7 dl = DataLoader(ds, batch_size=1, num_workers=1, shuffle=True)
----> 8 for x in dl:
      9     print(x)

~/.conda/envs/ml/lib/python3.6/site-packages/torch/utils/data/dataloader.py in __next__(self)
    635                 self.reorder_dict[idx] = batch
    636                 continue
--> 637             return self._process_next_batch(batch)
    638 
    639     next = __next__  # Python 2 compatibility

~/.conda/envs/ml/lib/python3.6/site-packages/torch/utils/data/dataloader.py in _process_next_batch(self, batch)
    656         self._put_indices()
    657         if isinstance(batch, ExceptionWrapper):
--> 658             raise batch.exc_type(batch.exc_msg)
    659         return batch
    660 

RuntimeError: Traceback (most recent call last):
  File "/home/swyoon/.conda/envs/ml/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 138, in _worker_loop
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/home/swyoon/.conda/envs/ml/lib/python3.6/site-packages/torch/utils/data/dataloader.py", line 138, in <listcomp>
    samples = collate_fn([dataset[i] for i in batch_indices])
  File "/home/swyoon/.conda/envs/ml/lib/python3.6/site-packages/torch/utils/data/dataset.py", line 40, in __getitem__
    return tuple(tensor[index] for tensor in self.tensors)
  File "/home/swyoon/.conda/envs/ml/lib/python3.6/site-packages/torch/utils/data/dataset.py", line 40, in <genexpr>
    return tuple(tensor[index] for tensor in self.tensors)
RuntimeError: CUDA error: initialization error

Is using a DataLoader when all my data is loaded on GPU?
Is the error an intended feature of Pytorch?