Repository: Lightning-AI/pytorch-lightning
File: pl_examples/basic_examples/dali_image_classifier.py (Python)

Code before:
    def define_graph(self):
        images, labels = self.source()
        return images, labels


class DALIClassificationLoader(DALIClassificationIterator):
    """
    This class extends DALI's original DALIClassificationIterator with the __len__() function so that we can call len() on it
    """

    def __init__(
            self,
            pipelines,
            size=-1,
            reader_name=None,
            auto_reset=False,
            last_batch_policy=LastBatchPolicy.FILL,
            dynamic_shape=False,
            last_batch_padded=False,
    ):
        super().__init__(pipelines, size, reader_name, auto_reset, dynamic_shape,
                         last_batch_policy=last_batch_policy, last_batch_padded=last_batch_padded)

    def __len__(self):
        batch_count = self._size // (self._num_gpus * self.batch_size)
        last_batch = 0 if self._last_batch_policy == LastBatchPolicy.DROP else 1
        return batch_count + last_batch


class LitClassifier(pl.LightningModule):
    def __init__(self, hidden_dim=128, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()

        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))
        x = torch.relu(self.l2(x))
        return x

    def split_batch(self, batch):
        return batch[0]["data"], batch[0]["label"].squeeze().long()

    def training_step(self, batch, batch_idx):
        x, y = self.split_batch(batch)
        y_hat = self(x)
        loss = F.cross_entropy(y_hat, y)
        return loss

Review comment (suggestion):
will this work also with older DALI versions, if not can we add a simple switch to something for `dali.__version__ > 0.28` and the old code for `<=0.27`

Code after:
    def __init__(self, batch_size, eii, num_threads, device_id):
        super(ExternalSourcePipeline, self).__init__(batch_size, num_threads, device_id, seed=12)
        self.source = ops.ExternalSource(source=eii, num_outputs=2)
        self.build()

    def define_graph(self):
        images, labels = self.source()
        return images, labels


class DALIClassificationLoader(DALIClassificationIterator):
    """
    This class extends DALI's original DALIClassificationIterator with the __len__() function so that we can call len() on it
    """

    def __init__(
            self,
            pipelines,
            size=-1,
            reader_name=None,
            auto_reset=False,
            fill_last_batch=True,
            dynamic_shape=False,
            last_batch_padded=False,
    ):
        if NEW_DALI_API:
            last_batch_policy = LastBatchPolicy.FILL if fill_last_batch else LastBatchPolicy.DROP
            super().__init__(pipelines, size, reader_name, auto_reset, dynamic_shape,
                             last_batch_policy=last_batch_policy, last_batch_padded=last_batch_padded)
        else:
            super().__init__(pipelines, size, reader_name, auto_reset, fill_last_batch,
                             dynamic_shape, last_batch_padded)
        self._fill_last_batch = fill_last_batch

    def __len__(self):
        batch_count = self._size // (self._num_gpus * self.batch_size)
        last_batch = 1 if self._fill_last_batch else 1
        return batch_count + last_batch


class LitClassifier(pl.LightningModule):
    def __init__(self, hidden_dim=128, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()

        self.l1 = torch.nn.Linear(28 * 28, self.hparams.hidden_dim)
        self.l2 = torch.nn.Linear(self.hparams.hidden_dim, 10)

    def forward(self, x):
        x = x.view(x.size(0), -1)
        x = torch.relu(self.l1(x))