Add device as arg to TensorDataset (#170)

Author: stes

cebra/data/datasets.py

@@ -73,8 +73,9 @@ def __init__(
         continuous: Union[torch.Tensor, npt.NDArray] = None,
         discrete: Union[torch.Tensor, npt.NDArray] = None,
         offset: int = 1,
+        device: str = "cpu"
     ):
-        super().__init__()
+        super().__init__(device=device)
         self.neural = self._to_tensor(neural, torch.FloatTensor).float()
         self.continuous = self._to_tensor(continuous, torch.FloatTensor)
         self.discrete = self._to_tensor(discrete, torch.LongTensor)