Dataset and module serve different purposes. I don’t think it makes sense to merge them together. Moreover, most of the time you want dataset to return CPU tensors and use multiprocessing dataloader and then convert the fetched data to GPU tensors before using them with the model. So I would suggest writing a custom wrapping class of a dataloader and a module, and when it is asked for a dataloader iterator, return something like
def iter_with_convert(self):
  for tensor in self.dataloader:
    yield tensor.to(self.device)