code: def test_auto_scale_batch_size_set_datamodule_attribute(
      tmpdir, use_hparams
  ):
      """Test that new batch size gets written to the correct hyperparameter attribute for datamodule."""
      hparams = {"batch_size": 2}
      before_batch_size = hparams["batch_size"]
  
      class HparamsBatchSizeDataModule(BoringDataModule):
          def __init__(self, batch_size):
              super().__init__()
              self.save_hyperparameters()
  
          def train_dataloader(self):
              return DataLoader(
                  self.random_train, batch_size=self.hparams.batch_size
              )
  
          def val_dataloader(self):
              return DataLoader(
                  RandomDataset(32, 64),
                  batch_size=self.hparams.batch_size,
              )
  
      datamodule_class = (
          HparamsBatchSizeDataModule
          if use_hparams
          else BatchSizeDataModule
      )
      datamodule = datamodule_class(batch_size=before_batch_size)
      model = BatchSizeModel(**hparams)
  
      trainer = Trainer(default_root_dir=tmpdir, max_epochs=1)
      tuner = Tuner(trainer)
      tuner.scale_batch_size(
          model, datamodule=datamodule, steps_per_trial=2, max_trials=4
      )
      after_batch_size = (
          datamodule.hparams.batch_size
          if use_hparams
          else datamodule.batch_size
      )
      assert trainer.datamodule == datamodule
      assert before_batch_size < after_batch_size
      assert after_batch_size <= len(trainer.train_dataloader.dataset)
  
imports: import logging
  import os
  from copy import deepcopy
  from unittest.mock import patch
  import pytest
  import torch
  from lightning.pytorch import Trainer
  from lightning.pytorch.callbacks.batch_size_finder import BatchSizeFinder
  from lightning.pytorch.demos.boring_classes import BoringDataModule, BoringModel, RandomDataset
  from lightning.pytorch.tuner.tuning import Tuner
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning_utilities.test.warning import no_warning_call
  from torch.utils.data import DataLoader
  from tests_pytorch.helpers.runif import RunIf
label_desc: Use image node_id 5 for calling a global function with example usage: test_auto_scale_batch_size_set_datamodule_attribute(tmpdir, use_hparams) and returns: DataLoader, DataLoader