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

Code before:
    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 DALI_VERSION >= '0.28.0':
            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))

Review comment (suggestion):
lets rather use `from distutils.version import LooseVersion`

Code after:
    This DALI pipeline class just contains the MNIST iterator
    """

    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)