import unittest
import torch
from torch.utils.data import TensorDataset, DataLoader, RandomSampler
import importlib
import io
import sys
import pandas as pd


def _dynamic(test, cond, success, fail):
    if cond:
        test._testMethodName = success
        test.assertTrue(True)
    else:
        test._testMethodName = fail
        test.fail(fail)


class TestUserCode(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.df = pd.read_csv(
            "https://content-media-cdn.codefinity.com/courses/1dd2b0f6-6ec0-40e6-a570-ed0ac2209666/section_2/iris.csv"
        )
        cls.features = torch.tensor(cls.df.drop(columns="species").values, dtype=torch.float32)
        cls.targets = torch.tensor(cls.df["species"].values, dtype=torch.long)

    # -------- Dataset exists --------
    def test_dataset_exists(self):
        import user_code
        cond = hasattr(user_code, "iris_dataset")
        _dynamic(
            self,
            cond,
            "`iris_dataset` is successfully created.",
            "`iris_dataset` is missing. Create it using `TensorDataset(features_tensor, target_tensor)`."
        )

    # -------- Dataset is TensorDataset --------
    def test_dataset_is_tensor_dataset(self):
        import user_code
        cond = isinstance(user_code.iris_dataset, TensorDataset)
        _dynamic(
            self,
            cond,
            "`iris_dataset` is a valid TensorDataset object.",
            "`iris_dataset` must be created using `TensorDataset(features_tensor, target_tensor)`."
        )

    # -------- Dataset content correctness --------
    def test_dataset_tensors_correct(self):
        import user_code
        ds = user_code.iris_dataset

        try:
            feats_ok = torch.allclose(ds.tensors[0], self.features)
            targets_ok = torch.allclose(ds.tensors[1], self.targets)
            cond = feats_ok and targets_ok
        except Exception:
            cond = False

        _dynamic(
            self,
            cond,
            "`iris_dataset` correctly stores the feature and target tensors.",
            "`iris_dataset` must contain `(features_tensor, target_tensor)` in that order."
        )

    # -------- DataLoader exists --------
    def test_dataloader_exists(self):
        import user_code
        cond = hasattr(user_code, "iris_loader")
        _dynamic(
            self,
            cond,
            "`iris_loader` is successfully created.",
            "`iris_loader` is missing. Wrap your dataset using `DataLoader(...)`."
        )

    # -------- DataLoader parameters --------
    def test_dataloader_params(self):
        import user_code
        loader = user_code.iris_loader

        cond = (
            isinstance(loader, DataLoader)
            and loader.batch_size == 32
            and isinstance(loader.sampler, RandomSampler)
        )

        _dynamic(
            self,
            cond,
            "`iris_loader` uses batch_size=32 and shuffles the data (RandomSampler).",
            "`iris_loader` must be created with `batch_size=32` and `shuffle=True` so that it uses a RandomSampler."
        )

    # -------- First batch correctness --------
    def test_first_batch_shape(self):
        import user_code

        loader = user_code.iris_loader
        batch_features, batch_targets = next(iter(loader))

        cond = (
            batch_features.ndim == 2
            and batch_features.shape[1] == 4
            and batch_targets.ndim == 1
        )

        _dynamic(
            self,
            cond,
            "Batches from `iris_loader` have the expected shapes.",
            "Batch shapes are incorrect. Features must be `[batch_size, 4]` and targets must be `[batch_size]`."
        )

    # -------- Output printing --------
    def test_output_printed(self):
        import user_code

        buf = io.StringIO()
        sys.stdout = buf
        importlib.reload(user_code)
        sys.stdout = sys.__stdout__

        cond = len(buf.getvalue().strip()) > 0

        _dynamic(
            self,
            cond,
            "The DataLoader output is printed successfully.",
            "No output detected. Ensure you print the first batch from the DataLoader."
        )


if __name__ == "__main__":
    unittest.main()


test_code.py

Transition from high-level libraries to the flexibility of PyTorch. Master the building blocks of deep learning by working with Tensors, the core data structure of PyTorch. Explore automatic differentiation with gradients, implement backpropagation, and build complete neural networks from scratch. Learn to handle real-world datasets and manage the full training cycle – from weight initialization to model evaluation.

Challenge: Preparing the Iris Dataset

Solution