From c10c4be600bc2d15a5260a166d34fea2d8389026 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Wed, 3 Apr 2024 15:42:12 +0100
Subject: [PATCH 01/11] Remove unnecessary import

---
 examples/advanced-pytorch/server.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/examples/advanced-pytorch/server.py b/examples/advanced-pytorch/server.py
index 489694ab1ea..d98a67d8288 100644
--- a/examples/advanced-pytorch/server.py
+++ b/examples/advanced-pytorch/server.py
@@ -10,8 +10,6 @@
 
 import warnings
 
-from flwr_datasets import FederatedDataset
-
 warnings.filterwarnings("ignore")
 
 

From 7dda131645f7c9444ab86a81a8dfb95e1ef02888 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Thu, 4 Apr 2024 15:00:52 +0100
Subject: [PATCH 02/11] Remove Poetry (version 1.7.1)

Usage:
  command [options] [arguments]

Options:
  -h, --help                 Display help for the given command. When no command is given display help for the list command.
  -q, --quiet                Do not output any message.
  -V, --version              Display this application version.
      --ansi                 Force ANSI output.
      --no-ansi              Disable ANSI output.
  -n, --no-interaction       Do not ask any interactive question.
      --no-plugins           Disables plugins.
      --no-cache             Disables Poetry source caches.
  -C, --directory=DIRECTORY  The working directory for the Poetry command (defaults to the current working directory).
  -v|vv|vvv, --verbose       Increase the verbosity of messages: 1 for normal output, 2 for more verbose output and 3 for debug.

Available commands:
  about              Shows information about Poetry.
  add                Adds a new dependency to pyproject.toml.
  build              Builds a package, as a tarball and a wheel by default.
  check              Validates the content of the pyproject.toml file and its consistency with the poetry.lock file.
  config             Manages configuration settings.
  export             Exports the lock file to alternative formats.
  help               Displays help for a command.
  init               Creates a basic pyproject.toml file in the current directory.
  install            Installs the project dependencies.
  list               Lists commands.
  lock               Locks the project dependencies.
  new                Creates a new Python project at <path>.
  publish            Publishes a package to a remote repository.
  remove             Removes a package from the project dependencies.
  run                Runs a command in the appropriate environment.
  search             Searches for packages on remote repositories.
  shell              Spawns a shell within the virtual environment.
  show               Shows information about packages.
  update             Update the dependencies as according to the pyproject.toml file.
  version            Shows the version of the project or bumps it when a valid bump rule is provided.

 cache
  cache clear        Clears a Poetry cache by name.
  cache list         List Poetry's caches.

 debug
  debug info         Shows debug information.
  debug resolve      Debugs dependency resolution.

 env
  env info           Displays information about the current environment.
  env list           Lists all virtualenvs associated with the current project.
  env remove         Remove virtual environments associated with the project.
  env use            Activates or creates a new virtualenv for the current project.

 self
  self add           Add additional packages to Poetry's runtime environment.
  self install       Install locked packages (incl. addons) required by this Poetry installation.
  self lock          Lock the Poetry installation's system requirements.
  self remove        Remove additional packages from Poetry's runtime environment.
  self show          Show packages from Poetry's runtime environment.
  self show plugins  Shows information about the currently installed plugins.
  self update        Updates Poetry to the latest version.

 source
  source add         Add source configuration for project.
  source remove      Remove source configured for the project.
  source show        Show information about sources configured for the project. and update
---
 examples/quickstart-pytorch/client.py        |   2 +-
 examples/quickstart-pytorch/client_0.py      | 152 +++++++++++++++++++
 examples/quickstart-pytorch/client_1.py      | 152 +++++++++++++++++++
 examples/quickstart-pytorch/pyproject.toml   |  26 ++--
 examples/quickstart-pytorch/requirements.txt |   5 -
 5 files changed, 320 insertions(+), 17 deletions(-)
 create mode 100644 examples/quickstart-pytorch/client_0.py
 create mode 100644 examples/quickstart-pytorch/client_1.py
 delete mode 100644 examples/quickstart-pytorch/requirements.txt

diff --git a/examples/quickstart-pytorch/client.py b/examples/quickstart-pytorch/client.py
index e58dbf7ea0b..2304857f661 100644
--- a/examples/quickstart-pytorch/client.py
+++ b/examples/quickstart-pytorch/client.py
@@ -107,7 +107,7 @@ def apply_transforms(batch):
 
 # Load model and data (simple CNN, CIFAR-10)
 net = Net().to(DEVICE)
-trainloader, testloader = load_data(partition_id=partition_id)
+trainloader, testloader = load_data(partition_id=0)
 
 
 # Define Flower client
diff --git a/examples/quickstart-pytorch/client_0.py b/examples/quickstart-pytorch/client_0.py
new file mode 100644
index 00000000000..2304857f661
--- /dev/null
+++ b/examples/quickstart-pytorch/client_0.py
@@ -0,0 +1,152 @@
+import argparse
+import warnings
+from collections import OrderedDict
+
+from flwr.client import NumPyClient, ClientApp
+from flwr_datasets import FederatedDataset
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision.transforms import Compose, Normalize, ToTensor
+from tqdm import tqdm
+
+
+# #############################################################################
+# 1. Regular PyTorch pipeline: nn.Module, train, test, and DataLoader
+# #############################################################################
+
+warnings.filterwarnings("ignore", category=UserWarning)
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class Net(nn.Module):
+    """Model (simple CNN adapted from 'PyTorch: A 60 Minute Blitz')"""
+
+    def __init__(self) -> None:
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 16 * 5 * 5)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return self.fc3(x)
+
+
+def train(net, trainloader, epochs):
+    """Train the model on the training set."""
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
+    for _ in range(epochs):
+        for batch in tqdm(trainloader, "Training"):
+            images = batch["img"]
+            labels = batch["label"]
+            optimizer.zero_grad()
+            criterion(net(images.to(DEVICE)), labels.to(DEVICE)).backward()
+            optimizer.step()
+
+
+def test(net, testloader):
+    """Validate the model on the test set."""
+    criterion = torch.nn.CrossEntropyLoss()
+    correct, loss = 0, 0.0
+    with torch.no_grad():
+        for batch in tqdm(testloader, "Testing"):
+            images = batch["img"].to(DEVICE)
+            labels = batch["label"].to(DEVICE)
+            outputs = net(images)
+            loss += criterion(outputs, labels).item()
+            correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
+    accuracy = correct / len(testloader.dataset)
+    return loss, accuracy
+
+
+def load_data(partition_id):
+    """Load partition CIFAR10 data."""
+    fds = FederatedDataset(dataset="cifar10", partitioners={"train": 3})
+    partition = fds.load_partition(partition_id)
+    # Divide data on each node: 80% train, 20% test
+    partition_train_test = partition.train_test_split(test_size=0.2)
+    pytorch_transforms = Compose(
+        [ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    )
+
+    def apply_transforms(batch):
+        """Apply transforms to the partition from FederatedDataset."""
+        batch["img"] = [pytorch_transforms(img) for img in batch["img"]]
+        return batch
+
+    partition_train_test = partition_train_test.with_transform(apply_transforms)
+    trainloader = DataLoader(partition_train_test["train"], batch_size=32, shuffle=True)
+    testloader = DataLoader(partition_train_test["test"], batch_size=32)
+    return trainloader, testloader
+
+
+# #############################################################################
+# 2. Federation of the pipeline with Flower
+# #############################################################################
+
+# Get partition id
+parser = argparse.ArgumentParser(description="Flower")
+parser.add_argument(
+    "--partition-id",
+    choices=[0, 1, 2],
+    default=0,
+    type=int,
+    help="Partition of the dataset divided into 3 iid partitions created artificially.",
+)
+partition_id = parser.parse_known_args()[0].partition_id
+
+# Load model and data (simple CNN, CIFAR-10)
+net = Net().to(DEVICE)
+trainloader, testloader = load_data(partition_id=0)
+
+
+# Define Flower client
+class FlowerClient(NumPyClient):
+    def get_parameters(self, config):
+        return [val.cpu().numpy() for _, val in net.state_dict().items()]
+
+    def set_parameters(self, parameters):
+        params_dict = zip(net.state_dict().keys(), parameters)
+        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
+        net.load_state_dict(state_dict, strict=True)
+
+    def fit(self, parameters, config):
+        self.set_parameters(parameters)
+        train(net, trainloader, epochs=1)
+        return self.get_parameters(config={}), len(trainloader.dataset), {}
+
+    def evaluate(self, parameters, config):
+        self.set_parameters(parameters)
+        loss, accuracy = test(net, testloader)
+        return loss, len(testloader.dataset), {"accuracy": accuracy}
+
+
+def client_fn(cid: str):
+    """Create and return an instance of Flower `Client`."""
+    return FlowerClient().to_client()
+
+
+# Flower ClientApp
+app = ClientApp(
+    client_fn=client_fn,
+)
+
+
+# Legacy mode
+if __name__ == "__main__":
+    from flwr.client import start_client
+
+    start_client(
+        server_address="127.0.0.1:8080",
+        client=FlowerClient().to_client(),
+    )
diff --git a/examples/quickstart-pytorch/client_1.py b/examples/quickstart-pytorch/client_1.py
new file mode 100644
index 00000000000..3b376fadafb
--- /dev/null
+++ b/examples/quickstart-pytorch/client_1.py
@@ -0,0 +1,152 @@
+import argparse
+import warnings
+from collections import OrderedDict
+
+from flwr.client import NumPyClient, ClientApp
+from flwr_datasets import FederatedDataset
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from torchvision.transforms import Compose, Normalize, ToTensor
+from tqdm import tqdm
+
+
+# #############################################################################
+# 1. Regular PyTorch pipeline: nn.Module, train, test, and DataLoader
+# #############################################################################
+
+warnings.filterwarnings("ignore", category=UserWarning)
+DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+class Net(nn.Module):
+    """Model (simple CNN adapted from 'PyTorch: A 60 Minute Blitz')"""
+
+    def __init__(self) -> None:
+        super(Net, self).__init__()
+        self.conv1 = nn.Conv2d(3, 6, 5)
+        self.pool = nn.MaxPool2d(2, 2)
+        self.conv2 = nn.Conv2d(6, 16, 5)
+        self.fc1 = nn.Linear(16 * 5 * 5, 120)
+        self.fc2 = nn.Linear(120, 84)
+        self.fc3 = nn.Linear(84, 10)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.pool(F.relu(self.conv1(x)))
+        x = self.pool(F.relu(self.conv2(x)))
+        x = x.view(-1, 16 * 5 * 5)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        return self.fc3(x)
+
+
+def train(net, trainloader, epochs):
+    """Train the model on the training set."""
+    criterion = torch.nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
+    for _ in range(epochs):
+        for batch in tqdm(trainloader, "Training"):
+            images = batch["img"]
+            labels = batch["label"]
+            optimizer.zero_grad()
+            criterion(net(images.to(DEVICE)), labels.to(DEVICE)).backward()
+            optimizer.step()
+
+
+def test(net, testloader):
+    """Validate the model on the test set."""
+    criterion = torch.nn.CrossEntropyLoss()
+    correct, loss = 0, 0.0
+    with torch.no_grad():
+        for batch in tqdm(testloader, "Testing"):
+            images = batch["img"].to(DEVICE)
+            labels = batch["label"].to(DEVICE)
+            outputs = net(images)
+            loss += criterion(outputs, labels).item()
+            correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
+    accuracy = correct / len(testloader.dataset)
+    return loss, accuracy
+
+
+def load_data(partition_id):
+    """Load partition CIFAR10 data."""
+    fds = FederatedDataset(dataset="cifar10", partitioners={"train": 3})
+    partition = fds.load_partition(partition_id)
+    # Divide data on each node: 80% train, 20% test
+    partition_train_test = partition.train_test_split(test_size=0.2)
+    pytorch_transforms = Compose(
+        [ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
+    )
+
+    def apply_transforms(batch):
+        """Apply transforms to the partition from FederatedDataset."""
+        batch["img"] = [pytorch_transforms(img) for img in batch["img"]]
+        return batch
+
+    partition_train_test = partition_train_test.with_transform(apply_transforms)
+    trainloader = DataLoader(partition_train_test["train"], batch_size=32, shuffle=True)
+    testloader = DataLoader(partition_train_test["test"], batch_size=32)
+    return trainloader, testloader
+
+
+# #############################################################################
+# 2. Federation of the pipeline with Flower
+# #############################################################################
+
+# Get partition id
+parser = argparse.ArgumentParser(description="Flower")
+parser.add_argument(
+    "--partition-id",
+    choices=[0, 1, 2],
+    default=0,
+    type=int,
+    help="Partition of the dataset divided into 3 iid partitions created artificially.",
+)
+partition_id = parser.parse_known_args()[0].partition_id
+
+# Load model and data (simple CNN, CIFAR-10)
+net = Net().to(DEVICE)
+trainloader, testloader = load_data(partition_id=1)
+
+
+# Define Flower client
+class FlowerClient(NumPyClient):
+    def get_parameters(self, config):
+        return [val.cpu().numpy() for _, val in net.state_dict().items()]
+
+    def set_parameters(self, parameters):
+        params_dict = zip(net.state_dict().keys(), parameters)
+        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
+        net.load_state_dict(state_dict, strict=True)
+
+    def fit(self, parameters, config):
+        self.set_parameters(parameters)
+        train(net, trainloader, epochs=1)
+        return self.get_parameters(config={}), len(trainloader.dataset), {}
+
+    def evaluate(self, parameters, config):
+        self.set_parameters(parameters)
+        loss, accuracy = test(net, testloader)
+        return loss, len(testloader.dataset), {"accuracy": accuracy}
+
+
+def client_fn(cid: str):
+    """Create and return an instance of Flower `Client`."""
+    return FlowerClient().to_client()
+
+
+# Flower ClientApp
+app = ClientApp(
+    client_fn=client_fn,
+)
+
+
+# Legacy mode
+if __name__ == "__main__":
+    from flwr.client import start_client
+
+    start_client(
+        server_address="127.0.0.1:8080",
+        client=FlowerClient().to_client(),
+    )
diff --git a/examples/quickstart-pytorch/pyproject.toml b/examples/quickstart-pytorch/pyproject.toml
index 7255e627471..ec49b5c412a 100644
--- a/examples/quickstart-pytorch/pyproject.toml
+++ b/examples/quickstart-pytorch/pyproject.toml
@@ -1,17 +1,21 @@
 [build-system]
-requires = ["poetry-core>=1.4.0"]
-build-backend = "poetry.core.masonry.api"
+requires = ["hatchling"]
+build-backend = "hatchling.build"
 
-[tool.poetry]
+[project]
 name = "quickstart-pytorch"
 version = "0.1.0"
 description = "PyTorch Federated Learning Quickstart with Flower"
-authors = ["The Flower Authors <hello@flower.ai>"]
+authors = [
+    { name = "The Flower Authors", email = "hello@flower.ai" },
+]
+dependencies = [
+    "flwr>=1.8.0,<2.0",
+    "flwr-datasets[vision]>=0.0.2,<1.0.0",
+    "torch==2.1.1",
+    "torchvision==0.16.1",
+    "tqdm==4.65.0"
+]
 
-[tool.poetry.dependencies]
-python = ">=3.8,<3.11"
-flwr = ">=1.8.0,<2.0"
-flwr-datasets = { extras = ["vision"], version = ">=0.0.2,<1.0.0" }
-torch = "2.1.1"
-torchvision = "0.16.1"
-tqdm = "4.65.0"
+[tool.hatch.build.targets.wheel]
+packages = ["."]
\ No newline at end of file
diff --git a/examples/quickstart-pytorch/requirements.txt b/examples/quickstart-pytorch/requirements.txt
deleted file mode 100644
index 12627809551..00000000000
--- a/examples/quickstart-pytorch/requirements.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-flwr>=1.8.0, <2.0
-flwr-datasets[vision]>=0.0.2, <1.0.0
-torch==2.1.1
-torchvision==0.16.1
-tqdm==4.65.0

From 188aa335337e1ae1f70f8d669095ca7741e1cde8 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 5 Apr 2024 10:21:35 +0100
Subject: [PATCH 03/11] Update readme

---
 examples/quickstart-pytorch/README.md | 98 +++++++++++++++------------
 examples/quickstart-pytorch/run.sh    | 17 -----
 2 files changed, 54 insertions(+), 61 deletions(-)
 delete mode 100755 examples/quickstart-pytorch/run.sh

diff --git a/examples/quickstart-pytorch/README.md b/examples/quickstart-pytorch/README.md
index 978191cc0ec..ae3c9d462fc 100644
--- a/examples/quickstart-pytorch/README.md
+++ b/examples/quickstart-pytorch/README.md
@@ -13,90 +13,100 @@ git clone --depth=1 https://github.com/adap/flower.git && mv flower/examples/qui
 This will create a new directory called `quickstart-pytorch` containing the following files:
 
 ```shell
--- pyproject.toml
--- requirements.txt
--- client.py
--- server.py
--- README.md
+.
+├── README.md
+├── client.py       # <-- for backwards compatibility
+├── client_0.py     # <-- contains `ClientApp`
+├── client_1.py     # <-- contains `ClientApp`
+├── pyproject.toml  # <-- dependencies
+└── server.py       # <-- contains `ServerApp`
 ```
 
 ### Installing Dependencies
 
-Project dependencies (such as `torch` and `flwr`) are defined in `pyproject.toml` and `requirements.txt`. We recommend [Poetry](https://python-poetry.org/docs/) to install those dependencies and manage your virtual environment ([Poetry installation](https://python-poetry.org/docs/#installation)) or [pip](https://pip.pypa.io/en/latest/development/), but feel free to use a different way of installing dependencies and managing virtual environments if you have other preferences.
+Project dependencies (such as `torch` and `flwr`) are defined in `pyproject.toml`. We recommend [venv](https://docs.python.org/3/library/venv.html) to manage your virtual environment ([Creating virtual environments](https://docs.python.org/3/library/venv.html#creating-virtual-environments)) and [pip](https://pip.pypa.io/en/latest/development/) to install those dependencies, but feel free to use a different way of installing dependencies and managing virtual environments if you have other preferences.
 
-#### Poetry
+#### `venv` and `pip`
 
 ```shell
-poetry install
-poetry shell
+python -m venv .venv
+source .venv/bin/activate
 ```
 
-Poetry will install all your dependencies in a newly created virtual environment. To verify that everything works correctly you can run the following command:
-
+Run the command below in your terminal to install the dependencies according to the `pyproject.toml`.
 ```shell
-poetry run python3 -c "import flwr"
+(.venv) pip install .
 ```
 
-If you don't see any errors you're good to go!
-
-#### pip
-
-Write the command below in your terminal to install the dependencies according to the configuration file requirements.txt.
+`pip` will install all your dependencies in the newly created and activated virtual environment. To verify that everything works correctly you can run the following command:
 
 ```shell
-pip install -r requirements.txt
+(.venv) python -c "import flwr"
 ```
 
+If you don't see any errors you're good to go!
+
 ______________________________________________________________________
 
-## Run Federated Learning with PyTorch and Flower
+## Run Federated Learning with PyTorch and `Flower Next`
 
-Afterwards you are ready to start the Flower server as well as the clients. You can simply start the server in a terminal as follows:
+Note that all of the following commands should be executed with your `.venv` virtual environment activated.
 
-```shell
-python3 server.py
+### 1. Start the long-running Flower server (SuperLink)
+
+```bash
+flower-superlink --insecure
 ```
 
-Now you are ready to start the Flower clients which will participate in the learning. We need to specify the partition id to
-use different partitions of the data on different nodes.  To do so simply open two more terminal windows and run the
-following commands.
+### 2. Start the long-running Flower clients (SuperNodes)
 
-Start client 1 in the first terminal:
+Start 2 Flower `SuperNodes` in 2 separate terminal windows, using:
 
-```shell
-python3 client.py --partition-id 0
+```bash
+flower-client-app client_0:app --insecure
+```
+and
+```bash
+flower-client-app client_1:app --insecure
 ```
 
-Start client 2 in the second terminal:
+### 3. Run the Flower App
 
-```shell
-python3 client.py --partition-id 1
+With both the long-running server (SuperLink) and two clients (SuperNode) up and running, we can now run the actual Flower App:
+
+```bash
+flower-server-app server:app --insecure
 ```
 
 You will see that PyTorch is starting a federated training. Look at the [code](https://github.com/adap/flower/tree/main/examples/quickstart-pytorch) for a detailed explanation.
 
 ______________________________________________________________________
 
-## Run Federated Learning with PyTorch and `Flower Next`
+## (Legacy) Run Federated Learning with PyTorch and Flower
 
-### 1. Start the long-running Flower server (SuperLink)
+> Note that the following steps refer to the legacy approach of deploying a Flower server and clients, and is no longer recommended. We recommend migrating to the Flower Next high-level API approach as announced [here](https://flower.ai/blog/2024-04-03-announcing-flower-1.8-release).
 
-```bash
-flower-superlink --insecure
+
+After installation, you are ready to start the Flower server as well as the clients. You can simply start the server in a terminal as follows:
+
+```shell
+python3 server.py
 ```
 
-### 2. Start the long-running Flower clients (SuperNodes)
+Now you are ready to start the Flower clients which will participate in the learning. We need to specify the partition id to
+use different partitions of the data on different nodes.  To do so simply open two more terminal windows and run the
+following commands.
 
-Start 2 Flower `SuperNodes` in 2 separate terminal windows, using:
+Start client 1 in the first terminal:
 
-```bash
-flower-client-app client:app --insecure
+```shell
+python3 client.py --partition-id 0
 ```
 
-### 3. Run the Flower App
-
-With both the long-running server (SuperLink) and two clients (SuperNode) up and running, we can now run the actual Flower App:
+Start client 2 in the second terminal:
 
-```bash
-flower-server-app server:app --insecure
+```shell
+python3 client.py --partition-id 1
 ```
+
+You will see that PyTorch is starting a federated training.
diff --git a/examples/quickstart-pytorch/run.sh b/examples/quickstart-pytorch/run.sh
deleted file mode 100755
index 6ca9c8cafec..00000000000
--- a/examples/quickstart-pytorch/run.sh
+++ /dev/null
@@ -1,17 +0,0 @@
-#!/bin/bash
-set -e
-cd "$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )"/
-
-echo "Starting server"
-python server.py &
-sleep 3  # Sleep for 3s to give the server enough time to start
-
-for i in $(seq 0 1); do
-    echo "Starting client $i"
-    python client.py --partition-id "$i" &
-done
-
-# Enable CTRL+C to stop all background processes
-trap "trap - SIGTERM && kill -- -$$" SIGINT SIGTERM
-# Wait for all background processes to complete
-wait

From 02287f7eee812c16bc5700d41d3e7b0f4dbae4fe Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 5 Apr 2024 10:55:07 +0100
Subject: [PATCH 04/11] Lint readme

---
 examples/quickstart-pytorch/README.md | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/quickstart-pytorch/README.md b/examples/quickstart-pytorch/README.md
index ae3c9d462fc..386149b1a92 100644
--- a/examples/quickstart-pytorch/README.md
+++ b/examples/quickstart-pytorch/README.md
@@ -34,6 +34,7 @@ source .venv/bin/activate
 ```
 
 Run the command below in your terminal to install the dependencies according to the `pyproject.toml`.
+
 ```shell
 (.venv) pip install .
 ```
@@ -65,7 +66,9 @@ Start 2 Flower `SuperNodes` in 2 separate terminal windows, using:
 ```bash
 flower-client-app client_0:app --insecure
 ```
+
 and
+
 ```bash
 flower-client-app client_1:app --insecure
 ```
@@ -86,7 +89,6 @@ ______________________________________________________________________
 
 > Note that the following steps refer to the legacy approach of deploying a Flower server and clients, and is no longer recommended. We recommend migrating to the Flower Next high-level API approach as announced [here](https://flower.ai/blog/2024-04-03-announcing-flower-1.8-release).
 
-
 After installation, you are ready to start the Flower server as well as the clients. You can simply start the server in a terminal as follows:
 
 ```shell

From 22469513e76457277b5d8c833735c28b1d612869 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 5 Apr 2024 11:15:19 +0100
Subject: [PATCH 05/11] Remove  and  import

---
 examples/quickstart-pytorch/client_0.py | 12 ------------
 examples/quickstart-pytorch/client_1.py | 12 ------------
 2 files changed, 24 deletions(-)

diff --git a/examples/quickstart-pytorch/client_0.py b/examples/quickstart-pytorch/client_0.py
index 2304857f661..412d7f93b7a 100644
--- a/examples/quickstart-pytorch/client_0.py
+++ b/examples/quickstart-pytorch/client_0.py
@@ -1,4 +1,3 @@
-import argparse
 import warnings
 from collections import OrderedDict
 
@@ -94,17 +93,6 @@ def apply_transforms(batch):
 # 2. Federation of the pipeline with Flower
 # #############################################################################
 
-# Get partition id
-parser = argparse.ArgumentParser(description="Flower")
-parser.add_argument(
-    "--partition-id",
-    choices=[0, 1, 2],
-    default=0,
-    type=int,
-    help="Partition of the dataset divided into 3 iid partitions created artificially.",
-)
-partition_id = parser.parse_known_args()[0].partition_id
-
 # Load model and data (simple CNN, CIFAR-10)
 net = Net().to(DEVICE)
 trainloader, testloader = load_data(partition_id=0)
diff --git a/examples/quickstart-pytorch/client_1.py b/examples/quickstart-pytorch/client_1.py
index 3b376fadafb..4221150bb68 100644
--- a/examples/quickstart-pytorch/client_1.py
+++ b/examples/quickstart-pytorch/client_1.py
@@ -1,4 +1,3 @@
-import argparse
 import warnings
 from collections import OrderedDict
 
@@ -94,17 +93,6 @@ def apply_transforms(batch):
 # 2. Federation of the pipeline with Flower
 # #############################################################################
 
-# Get partition id
-parser = argparse.ArgumentParser(description="Flower")
-parser.add_argument(
-    "--partition-id",
-    choices=[0, 1, 2],
-    default=0,
-    type=int,
-    help="Partition of the dataset divided into 3 iid partitions created artificially.",
-)
-partition_id = parser.parse_known_args()[0].partition_id
-
 # Load model and data (simple CNN, CIFAR-10)
 net = Net().to(DEVICE)
 trainloader, testloader = load_data(partition_id=1)

From b555b087d460995e38145ecde82bffdd7bbd6ca8 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 5 Apr 2024 13:18:53 +0100
Subject: [PATCH 06/11] Tidy up environment guide

---
 examples/quickstart-pytorch/README.md | 20 +++++---------------
 1 file changed, 5 insertions(+), 15 deletions(-)

diff --git a/examples/quickstart-pytorch/README.md b/examples/quickstart-pytorch/README.md
index 386149b1a92..83947a66a8c 100644
--- a/examples/quickstart-pytorch/README.md
+++ b/examples/quickstart-pytorch/README.md
@@ -24,25 +24,17 @@ This will create a new directory called `quickstart-pytorch` containing the foll
 
 ### Installing Dependencies
 
-Project dependencies (such as `torch` and `flwr`) are defined in `pyproject.toml`. We recommend [venv](https://docs.python.org/3/library/venv.html) to manage your virtual environment ([Creating virtual environments](https://docs.python.org/3/library/venv.html#creating-virtual-environments)) and [pip](https://pip.pypa.io/en/latest/development/) to install those dependencies, but feel free to use a different way of installing dependencies and managing virtual environments if you have other preferences.
-
-#### `venv` and `pip`
-
-```shell
-python -m venv .venv
-source .venv/bin/activate
-```
-
-Run the command below in your terminal to install the dependencies according to the `pyproject.toml`.
+Project dependencies (such as `torch` and `flwr`) are defined in `pyproject.toml`. You can install the dependencies by invoking `pip`: 
 
 ```shell
-(.venv) pip install .
+# From a new python environment, run:
+pip install .
 ```
 
-`pip` will install all your dependencies in the newly created and activated virtual environment. To verify that everything works correctly you can run the following command:
+Then, to verify that everything works correctly you can run the following command:
 
 ```shell
-(.venv) python -c "import flwr"
+python -c "import flwr"
 ```
 
 If you don't see any errors you're good to go!
@@ -51,8 +43,6 @@ ______________________________________________________________________
 
 ## Run Federated Learning with PyTorch and `Flower Next`
 
-Note that all of the following commands should be executed with your `.venv` virtual environment activated.
-
 ### 1. Start the long-running Flower server (SuperLink)
 
 ```bash

From 6d746f2ad8644bb96bbc6e78d131023df43eb4da Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 5 Apr 2024 14:13:32 +0100
Subject: [PATCH 07/11] Revert client split and readme

---
 examples/quickstart-pytorch/README.md   |  66 +++++------
 examples/quickstart-pytorch/client_0.py | 140 ------------------------
 examples/quickstart-pytorch/client_1.py | 140 ------------------------
 3 files changed, 27 insertions(+), 319 deletions(-)
 delete mode 100644 examples/quickstart-pytorch/client_0.py
 delete mode 100644 examples/quickstart-pytorch/client_1.py

diff --git a/examples/quickstart-pytorch/README.md b/examples/quickstart-pytorch/README.md
index 83947a66a8c..174b7fb1e8d 100644
--- a/examples/quickstart-pytorch/README.md
+++ b/examples/quickstart-pytorch/README.md
@@ -15,9 +15,7 @@ This will create a new directory called `quickstart-pytorch` containing the foll
 ```shell
 .
 ├── README.md
-├── client.py       # <-- for backwards compatibility
-├── client_0.py     # <-- contains `ClientApp`
-├── client_1.py     # <-- contains `ClientApp`
+├── client.py       # <-- contains `ClientApp`
 ├── pyproject.toml  # <-- dependencies
 └── server.py       # <-- contains `ServerApp`
 ```
@@ -41,64 +39,54 @@ If you don't see any errors you're good to go!
 
 ______________________________________________________________________
 
-## Run Federated Learning with PyTorch and `Flower Next`
+## Run Federated Learning with PyTorch and Flower
 
-### 1. Start the long-running Flower server (SuperLink)
+Afterwards you are ready to start the Flower server as well as the clients. You can simply start the server in a terminal as follows:
 
-```bash
-flower-superlink --insecure
+```shell
+python3 server.py
 ```
 
-### 2. Start the long-running Flower clients (SuperNodes)
-
-Start 2 Flower `SuperNodes` in 2 separate terminal windows, using:
-
-```bash
-flower-client-app client_0:app --insecure
-```
+Now you are ready to start the Flower clients which will participate in the learning. We need to specify the partition id to
+use different partitions of the data on different nodes.  To do so simply open two more terminal windows and run the
+following commands.
 
-and
+Start client 1 in the first terminal:
 
-```bash
-flower-client-app client_1:app --insecure
+```shell
+python3 client.py --partition-id 0
 ```
 
-### 3. Run the Flower App
-
-With both the long-running server (SuperLink) and two clients (SuperNode) up and running, we can now run the actual Flower App:
+Start client 2 in the second terminal:
 
-```bash
-flower-server-app server:app --insecure
+```shell
+python3 client.py --partition-id 1
 ```
 
 You will see that PyTorch is starting a federated training. Look at the [code](https://github.com/adap/flower/tree/main/examples/quickstart-pytorch) for a detailed explanation.
 
 ______________________________________________________________________
 
-## (Legacy) Run Federated Learning with PyTorch and Flower
-
-> Note that the following steps refer to the legacy approach of deploying a Flower server and clients, and is no longer recommended. We recommend migrating to the Flower Next high-level API approach as announced [here](https://flower.ai/blog/2024-04-03-announcing-flower-1.8-release).
+## Run Federated Learning with PyTorch and `Flower Next`
 
-After installation, you are ready to start the Flower server as well as the clients. You can simply start the server in a terminal as follows:
+### 1. Start the long-running Flower server (SuperLink)
 
-```shell
-python3 server.py
+```bash
+flower-superlink --insecure
 ```
 
-Now you are ready to start the Flower clients which will participate in the learning. We need to specify the partition id to
-use different partitions of the data on different nodes.  To do so simply open two more terminal windows and run the
-following commands.
+### 2. Start the long-running Flower clients (SuperNodes)
 
-Start client 1 in the first terminal:
+Start 2 Flower `SuperNodes` in 2 separate terminal windows, using:
 
-```shell
-python3 client.py --partition-id 0
+```bash
+flower-client-app client:app --insecure
 ```
 
-Start client 2 in the second terminal:
+### 3. Run the Flower App
 
-```shell
-python3 client.py --partition-id 1
-```
+With both the long-running server (SuperLink) and two clients (SuperNode) up and running, we can now run the actual Flower App:
 
-You will see that PyTorch is starting a federated training.
+```bash
+flower-server-app server:app --insecure
+```
\ No newline at end of file
diff --git a/examples/quickstart-pytorch/client_0.py b/examples/quickstart-pytorch/client_0.py
deleted file mode 100644
index 412d7f93b7a..00000000000
--- a/examples/quickstart-pytorch/client_0.py
+++ /dev/null
@@ -1,140 +0,0 @@
-import warnings
-from collections import OrderedDict
-
-from flwr.client import NumPyClient, ClientApp
-from flwr_datasets import FederatedDataset
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.utils.data import DataLoader
-from torchvision.transforms import Compose, Normalize, ToTensor
-from tqdm import tqdm
-
-
-# #############################################################################
-# 1. Regular PyTorch pipeline: nn.Module, train, test, and DataLoader
-# #############################################################################
-
-warnings.filterwarnings("ignore", category=UserWarning)
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
-class Net(nn.Module):
-    """Model (simple CNN adapted from 'PyTorch: A 60 Minute Blitz')"""
-
-    def __init__(self) -> None:
-        super(Net, self).__init__()
-        self.conv1 = nn.Conv2d(3, 6, 5)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(6, 16, 5)
-        self.fc1 = nn.Linear(16 * 5 * 5, 120)
-        self.fc2 = nn.Linear(120, 84)
-        self.fc3 = nn.Linear(84, 10)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.pool(F.relu(self.conv1(x)))
-        x = self.pool(F.relu(self.conv2(x)))
-        x = x.view(-1, 16 * 5 * 5)
-        x = F.relu(self.fc1(x))
-        x = F.relu(self.fc2(x))
-        return self.fc3(x)
-
-
-def train(net, trainloader, epochs):
-    """Train the model on the training set."""
-    criterion = torch.nn.CrossEntropyLoss()
-    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
-    for _ in range(epochs):
-        for batch in tqdm(trainloader, "Training"):
-            images = batch["img"]
-            labels = batch["label"]
-            optimizer.zero_grad()
-            criterion(net(images.to(DEVICE)), labels.to(DEVICE)).backward()
-            optimizer.step()
-
-
-def test(net, testloader):
-    """Validate the model on the test set."""
-    criterion = torch.nn.CrossEntropyLoss()
-    correct, loss = 0, 0.0
-    with torch.no_grad():
-        for batch in tqdm(testloader, "Testing"):
-            images = batch["img"].to(DEVICE)
-            labels = batch["label"].to(DEVICE)
-            outputs = net(images)
-            loss += criterion(outputs, labels).item()
-            correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
-    accuracy = correct / len(testloader.dataset)
-    return loss, accuracy
-
-
-def load_data(partition_id):
-    """Load partition CIFAR10 data."""
-    fds = FederatedDataset(dataset="cifar10", partitioners={"train": 3})
-    partition = fds.load_partition(partition_id)
-    # Divide data on each node: 80% train, 20% test
-    partition_train_test = partition.train_test_split(test_size=0.2)
-    pytorch_transforms = Compose(
-        [ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
-    )
-
-    def apply_transforms(batch):
-        """Apply transforms to the partition from FederatedDataset."""
-        batch["img"] = [pytorch_transforms(img) for img in batch["img"]]
-        return batch
-
-    partition_train_test = partition_train_test.with_transform(apply_transforms)
-    trainloader = DataLoader(partition_train_test["train"], batch_size=32, shuffle=True)
-    testloader = DataLoader(partition_train_test["test"], batch_size=32)
-    return trainloader, testloader
-
-
-# #############################################################################
-# 2. Federation of the pipeline with Flower
-# #############################################################################
-
-# Load model and data (simple CNN, CIFAR-10)
-net = Net().to(DEVICE)
-trainloader, testloader = load_data(partition_id=0)
-
-
-# Define Flower client
-class FlowerClient(NumPyClient):
-    def get_parameters(self, config):
-        return [val.cpu().numpy() for _, val in net.state_dict().items()]
-
-    def set_parameters(self, parameters):
-        params_dict = zip(net.state_dict().keys(), parameters)
-        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
-        net.load_state_dict(state_dict, strict=True)
-
-    def fit(self, parameters, config):
-        self.set_parameters(parameters)
-        train(net, trainloader, epochs=1)
-        return self.get_parameters(config={}), len(trainloader.dataset), {}
-
-    def evaluate(self, parameters, config):
-        self.set_parameters(parameters)
-        loss, accuracy = test(net, testloader)
-        return loss, len(testloader.dataset), {"accuracy": accuracy}
-
-
-def client_fn(cid: str):
-    """Create and return an instance of Flower `Client`."""
-    return FlowerClient().to_client()
-
-
-# Flower ClientApp
-app = ClientApp(
-    client_fn=client_fn,
-)
-
-
-# Legacy mode
-if __name__ == "__main__":
-    from flwr.client import start_client
-
-    start_client(
-        server_address="127.0.0.1:8080",
-        client=FlowerClient().to_client(),
-    )
diff --git a/examples/quickstart-pytorch/client_1.py b/examples/quickstart-pytorch/client_1.py
deleted file mode 100644
index 4221150bb68..00000000000
--- a/examples/quickstart-pytorch/client_1.py
+++ /dev/null
@@ -1,140 +0,0 @@
-import warnings
-from collections import OrderedDict
-
-from flwr.client import NumPyClient, ClientApp
-from flwr_datasets import FederatedDataset
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.utils.data import DataLoader
-from torchvision.transforms import Compose, Normalize, ToTensor
-from tqdm import tqdm
-
-
-# #############################################################################
-# 1. Regular PyTorch pipeline: nn.Module, train, test, and DataLoader
-# #############################################################################
-
-warnings.filterwarnings("ignore", category=UserWarning)
-DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
-class Net(nn.Module):
-    """Model (simple CNN adapted from 'PyTorch: A 60 Minute Blitz')"""
-
-    def __init__(self) -> None:
-        super(Net, self).__init__()
-        self.conv1 = nn.Conv2d(3, 6, 5)
-        self.pool = nn.MaxPool2d(2, 2)
-        self.conv2 = nn.Conv2d(6, 16, 5)
-        self.fc1 = nn.Linear(16 * 5 * 5, 120)
-        self.fc2 = nn.Linear(120, 84)
-        self.fc3 = nn.Linear(84, 10)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.pool(F.relu(self.conv1(x)))
-        x = self.pool(F.relu(self.conv2(x)))
-        x = x.view(-1, 16 * 5 * 5)
-        x = F.relu(self.fc1(x))
-        x = F.relu(self.fc2(x))
-        return self.fc3(x)
-
-
-def train(net, trainloader, epochs):
-    """Train the model on the training set."""
-    criterion = torch.nn.CrossEntropyLoss()
-    optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
-    for _ in range(epochs):
-        for batch in tqdm(trainloader, "Training"):
-            images = batch["img"]
-            labels = batch["label"]
-            optimizer.zero_grad()
-            criterion(net(images.to(DEVICE)), labels.to(DEVICE)).backward()
-            optimizer.step()
-
-
-def test(net, testloader):
-    """Validate the model on the test set."""
-    criterion = torch.nn.CrossEntropyLoss()
-    correct, loss = 0, 0.0
-    with torch.no_grad():
-        for batch in tqdm(testloader, "Testing"):
-            images = batch["img"].to(DEVICE)
-            labels = batch["label"].to(DEVICE)
-            outputs = net(images)
-            loss += criterion(outputs, labels).item()
-            correct += (torch.max(outputs.data, 1)[1] == labels).sum().item()
-    accuracy = correct / len(testloader.dataset)
-    return loss, accuracy
-
-
-def load_data(partition_id):
-    """Load partition CIFAR10 data."""
-    fds = FederatedDataset(dataset="cifar10", partitioners={"train": 3})
-    partition = fds.load_partition(partition_id)
-    # Divide data on each node: 80% train, 20% test
-    partition_train_test = partition.train_test_split(test_size=0.2)
-    pytorch_transforms = Compose(
-        [ToTensor(), Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))]
-    )
-
-    def apply_transforms(batch):
-        """Apply transforms to the partition from FederatedDataset."""
-        batch["img"] = [pytorch_transforms(img) for img in batch["img"]]
-        return batch
-
-    partition_train_test = partition_train_test.with_transform(apply_transforms)
-    trainloader = DataLoader(partition_train_test["train"], batch_size=32, shuffle=True)
-    testloader = DataLoader(partition_train_test["test"], batch_size=32)
-    return trainloader, testloader
-
-
-# #############################################################################
-# 2. Federation of the pipeline with Flower
-# #############################################################################
-
-# Load model and data (simple CNN, CIFAR-10)
-net = Net().to(DEVICE)
-trainloader, testloader = load_data(partition_id=1)
-
-
-# Define Flower client
-class FlowerClient(NumPyClient):
-    def get_parameters(self, config):
-        return [val.cpu().numpy() for _, val in net.state_dict().items()]
-
-    def set_parameters(self, parameters):
-        params_dict = zip(net.state_dict().keys(), parameters)
-        state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
-        net.load_state_dict(state_dict, strict=True)
-
-    def fit(self, parameters, config):
-        self.set_parameters(parameters)
-        train(net, trainloader, epochs=1)
-        return self.get_parameters(config={}), len(trainloader.dataset), {}
-
-    def evaluate(self, parameters, config):
-        self.set_parameters(parameters)
-        loss, accuracy = test(net, testloader)
-        return loss, len(testloader.dataset), {"accuracy": accuracy}
-
-
-def client_fn(cid: str):
-    """Create and return an instance of Flower `Client`."""
-    return FlowerClient().to_client()
-
-
-# Flower ClientApp
-app = ClientApp(
-    client_fn=client_fn,
-)
-
-
-# Legacy mode
-if __name__ == "__main__":
-    from flwr.client import start_client
-
-    start_client(
-        server_address="127.0.0.1:8080",
-        client=FlowerClient().to_client(),
-    )

From 6cb764bca85f87c02ce75d92bc380330c29d8269 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 5 Apr 2024 14:13:48 +0100
Subject: [PATCH 08/11] Lint readme

---
 examples/quickstart-pytorch/README.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/quickstart-pytorch/README.md b/examples/quickstart-pytorch/README.md
index 174b7fb1e8d..c89b2d5c47d 100644
--- a/examples/quickstart-pytorch/README.md
+++ b/examples/quickstart-pytorch/README.md
@@ -22,7 +22,7 @@ This will create a new directory called `quickstart-pytorch` containing the foll
 
 ### Installing Dependencies
 
-Project dependencies (such as `torch` and `flwr`) are defined in `pyproject.toml`. You can install the dependencies by invoking `pip`: 
+Project dependencies (such as `torch` and `flwr`) are defined in `pyproject.toml`. You can install the dependencies by invoking `pip`:
 
 ```shell
 # From a new python environment, run:
@@ -89,4 +89,4 @@ With both the long-running server (SuperLink) and two clients (SuperNode) up and
 
 ```bash
 flower-server-app server:app --insecure
-```
\ No newline at end of file
+```

From d723fa2726e533f5710a2986b1374282067d7bf7 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 5 Apr 2024 14:18:35 +0100
Subject: [PATCH 09/11] Revert more changes

---
 examples/quickstart-pytorch/README.md | 9 ++++-----
 examples/quickstart-pytorch/client.py | 2 +-
 2 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/examples/quickstart-pytorch/README.md b/examples/quickstart-pytorch/README.md
index c89b2d5c47d..8dc040d5024 100644
--- a/examples/quickstart-pytorch/README.md
+++ b/examples/quickstart-pytorch/README.md
@@ -13,11 +13,10 @@ git clone --depth=1 https://github.com/adap/flower.git && mv flower/examples/qui
 This will create a new directory called `quickstart-pytorch` containing the following files:
 
 ```shell
-.
-├── README.md
-├── client.py       # <-- contains `ClientApp`
-├── pyproject.toml  # <-- dependencies
-└── server.py       # <-- contains `ServerApp`
+-- pyproject.toml
+-- client.py
+-- server.py
+-- README.md
 ```
 
 ### Installing Dependencies
diff --git a/examples/quickstart-pytorch/client.py b/examples/quickstart-pytorch/client.py
index b9eb9a1c703..be4be88b8f8 100644
--- a/examples/quickstart-pytorch/client.py
+++ b/examples/quickstart-pytorch/client.py
@@ -107,7 +107,7 @@ def apply_transforms(batch):
 
 # Load model and data (simple CNN, CIFAR-10)
 net = Net().to(DEVICE)
-trainloader, testloader = load_data(partition_id=0)
+trainloader, testloader = load_data(partition_id=partition_id)
 
 
 # Define Flower client

From be491bf005cd786fca46c45e869768b617790d14 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 19 Apr 2024 08:10:47 +0100
Subject: [PATCH 10/11] Address comments

---
 examples/advanced-pytorch/server.py        | 2 ++
 examples/quickstart-pytorch/pyproject.toml | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/examples/advanced-pytorch/server.py b/examples/advanced-pytorch/server.py
index d98a67d8288..489694ab1ea 100644
--- a/examples/advanced-pytorch/server.py
+++ b/examples/advanced-pytorch/server.py
@@ -10,6 +10,8 @@
 
 import warnings
 
+from flwr_datasets import FederatedDataset
+
 warnings.filterwarnings("ignore")
 
 
diff --git a/examples/quickstart-pytorch/pyproject.toml b/examples/quickstart-pytorch/pyproject.toml
index ec49b5c412a..4692958d449 100644
--- a/examples/quickstart-pytorch/pyproject.toml
+++ b/examples/quickstart-pytorch/pyproject.toml
@@ -18,4 +18,4 @@ dependencies = [
 ]
 
 [tool.hatch.build.targets.wheel]
-packages = ["."]
\ No newline at end of file
+packages = ["."]

From a08e68c34d256ffb081e3ee63b63e74dce2c9266 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Fri, 19 Apr 2024 08:13:05 +0100
Subject: [PATCH 11/11] Change from python to python3 for clarity

---
 examples/quickstart-pytorch/README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/quickstart-pytorch/README.md b/examples/quickstart-pytorch/README.md
index 8dc040d5024..93d6a593f36 100644
--- a/examples/quickstart-pytorch/README.md
+++ b/examples/quickstart-pytorch/README.md
@@ -31,7 +31,7 @@ pip install .
 Then, to verify that everything works correctly you can run the following command:
 
 ```shell
-python -c "import flwr"
+python3 -c "import flwr"
 ```
 
 If you don't see any errors you're good to go!