old_contents: import pytorch_lightning as pl
  import torch
  
  
  class DCLightningDatasetBatch:
  
    def __init__(self, batch):
      X = [batch[0]]
      y = [batch[1]]
      w = [batch[2]]
      self.batch_list = [X, y, w]
  
  
  def collate_dataset_wrapper(batch):
    return DCLightningDatasetBatch(batch)
  
  
  class DCLightningDatasetModule(pl.LightningDataModule):
  
    def __init__(self, dataset, batch_size, collate_fn):
      super().__init__()
      self._batch_size = batch_size
      self._dataset = dataset.make_pytorch_dataset(batch_size=self._batch_size)
      self.collate_fn = collate_fn
  
    def setup(self, stage):
      self.train_dataset = self._dataset
  
    def train_dataloader(self):
      return torch.utils.data.DataLoader(
          self.train_dataset,
          batch_size=None,
          collate_fn=self.collate_fn,
          shuffle=False,
      )
  
new_contents: import pytorch_lightning as pl
  import torch
  
  
  class DCLightningDatasetBatch:
  
    def __init__(self, batch):
      X = [batch[0]]
      y = [batch[1]]
      w = [batch[2]]
      self.batch_list = [X, y, w]
  
  
  def collate_dataset_wrapper(batch):
    return DCLightningDatasetBatch(batch)
  
  
  class DCLightningDatasetModule(pl.LightningDataModule):
    """DeepChem Lightning Dataset Module to be used with the DCLightningModule and a Lightning trainer.
  
    This module wraps over the the deepchem pytorch dataset and dataloader providing a generic interface to run training.
    https://pytorch-lightning.readthedocs.io/en/stable/extensions/datamodules.html
  
    Notes
    -----
    This class requires PyTorch to be installed.
    """
  
    def __init__(self, dataset, batch_size, collate_fn):
      """Create a new DCLightningDatasetModule.
  
      Parameters
      ----------
      dataset: A deepchem dataset.
      batch_size: Batch size for the dataloader.
      collate_fn: Method to collate instances across batch.
      """
      super().__init__()
      self._batch_size = batch_size
      self._dataset = dataset.make_pytorch_dataset(batch_size=self._batch_size)
      self.collate_fn = collate_fn
  
    def setup(self, stage):
      self.train_dataset = self._dataset
  
    def train_dataloader(self):
      """Returns the train dataloader from train dataset.
  
      Returns
      -------
      dataloader: train dataloader to be used with DCLightningModule.
      """
      dataloader = torch.utils.data.DataLoader(
          self.train_dataset,
          batch_size=None,
          collate_fn=self.collate_fn,
          shuffle=False,
      )
      return dataloader
  
subject: Add docstring