Fix pylint errors

adap · Nov 10, 2023 · 4316466 · 4316466
1 parent d419afc
commit 4316466
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/datasets/flwr_datasets/partitioner/iid_partitioner_test.py b/datasets/flwr_datasets/partitioner/iid_partitioner_test.py
@@ -18,7 +18,6 @@
 import unittest
 from typing import Tuple
 
-import numpy as np
 from parameterized import parameterized
 
 from datasets import Dataset
@@ -107,10 +106,10 @@ def test_load_partition_correct_data(
         partition = partitioner.load_partition(partition_index)
         row_id = 0
         self.assertEqual(
-            partition["features"][row_id],
+            partition[row_id]["features"],
             # Note it's contiguous so partition_size * partition_index gets the first
             # element of the partition of partition_index
-            dataset["features"][partition_size * partition_index + row_id],
+            dataset[partition_size * partition_index + row_id]["features"],
         )
 
     @parameterized.expand(  # type: ignore