From 6021a0f46ebfec4c8f20bef6915ca84c86849627 Mon Sep 17 00:00:00 2001
From: Chong Shen Ng <chong.shen@flower.ai>
Date: Wed, 6 Nov 2024 11:44:08 +0000
Subject: [PATCH] Add flwr_datasets as known_first_party

---
 e2e/docker/client.py | 2 +-
 pyproject.toml       | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/e2e/docker/client.py b/e2e/docker/client.py
index 44313c7c3af6..98a343f239f3 100644
--- a/e2e/docker/client.py
+++ b/e2e/docker/client.py
@@ -4,12 +4,12 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from flwr_datasets import FederatedDataset
 from torch.utils.data import DataLoader, Subset
 from torchvision.transforms import Compose, Normalize, ToTensor
 
 from flwr.client import ClientApp, NumPyClient
 from flwr.common import Context
+from flwr_datasets import FederatedDataset
 
 # #############################################################################
 # 1. Regular PyTorch pipeline: nn.Module, train, test, and DataLoader
diff --git a/pyproject.toml b/pyproject.toml
index a3887bc0293c..88aa2ebad211 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -147,7 +147,7 @@ extend_exclude = [
 
 [tool.isort]
 profile = "black"
-known_first_party = ["flwr", "flwr_tool"]
+known_first_party = ["flwr", "flwr_tool", "flwr_datasets"]
 
 [tool.black]
 line-length = 88