code: def __init__(self, *args, ds_type="map"):
      self.ds_type = ds_type
  
      if ds_type == "map":
          self.ds = torch.utils.data.TensorDataset(*args)
      else:
          self.ds = DummyIterableDataset(*args)
  
file_path: TorchDrift/test/test_utils.py
imports: import pytest
  import torchdrift
  import torch
  import torch.utils.data
label_desc: Use image node_id 1 to create a new TensorDataModule object  with example: obj = TensorDataModule()