class_docstr: >>> ImageNetLightningModel(data_path='missing')  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
  ImageNetLightningModel(
    (model): ResNet(...)
  )
class_docstr_tok: [">",">",">","ImageNetLightningModel","(","data_path='missing","'",")","#","doctest",":","+ELLIPSIS","+NORMALIZE_WHITESPACE","ImageNetLightningModel","(","(","model",")",":","ResNet","(","...",")",")"]
code: def setup(self, stage: str):
      if isinstance(self.trainer.strategy, ParallelStrategy):
          # When using a single GPU per process and per `DistributedDataParallel`, we need to divide the batch size
          # ourselves based on the total number of GPUs we have
          num_processes = max(1, self.trainer.strategy.num_processes)
          self.batch_size = int(self.batch_size / num_processes)
          self.workers = int(self.workers / num_processes)
  
      if stage == "fit":
          normalize = transforms.Normalize(
              mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
          )
          train_dir = os.path.join(self.data_path, "train")
          self.train_dataset = datasets.ImageFolder(
              train_dir,
              transforms.Compose(
                  [
                      transforms.RandomResizedCrop(224),
                      transforms.RandomHorizontalFlip(),
                      transforms.ToTensor(),
                      normalize,
                  ]
              ),
          )
      # all stages will use the eval dataset
      normalize = transforms.Normalize(
          mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
      )
      val_dir = os.path.join(self.data_path, "val")
      self.eval_dataset = datasets.ImageFolder(
          val_dir,
          transforms.Compose(
              [
                  transforms.Resize(256),
                  transforms.CenterCrop(224),
                  transforms.ToTensor(),
                  normalize,
              ]
          ),
      )
  
imports: import os
  from typing import Optional
  import torch
  import torch.nn.functional
  import torch.nn.parallel
  import torch.optim
  import torch.optim.lr_scheduler
  import torch.utils.data
  import torch.utils.data.distributed
  import torchvision.datasets
  import torchvision.transforms
  from lightning.pytorch import LightningModule
  from lightning.pytorch.callbacks import ModelCheckpoint, TQDMProgressBar
  from lightning.pytorch.cli import LightningCLI
  from lightning.pytorch.strategies import ParallelStrategy
  from lightning.pytorch.utilities.model_helpers import get_torchvision_model
  from torch.utils.data import Dataset
  from torchmetrics import Accuracy
label_desc: Use image node_id 8 for calling the ImageNetLightningModel obj's underlying member method code with example usage: obj.setup(stage) without return types