diff --git a/baselines/fedht/fedht/aggregate.py b/baselines/fedht/fedht/aggregate.py
index b6ef28718e3..43dceaf9415 100644
--- a/baselines/fedht/fedht/aggregate.py
+++ b/baselines/fedht/fedht/aggregate.py
@@ -416,7 +416,6 @@ def aggregate_hardthreshold(
     if num_keep <= 0:
         raise ValueError("k must be a positive integer.")
 
-    """Compute weighted average."""
     # Calculate the total number of examples used during training
     num_examples_total = sum(num_examples for (_, num_examples) in results)
 
@@ -438,10 +437,10 @@ def aggregate_hardthreshold(
         # ignoring second element in results[i] skips over number of observations
         # j: iterates through all layers of a model
         # k: iterates through the slices of each layer
-        for i in range(len(results)):
-            for j in range(len(results[i][0])):
-                for k in range(len(results[i][0][j])):
-                    results[i][0][j][k] = hardthreshold(results[i][0][j][k], num_keep)
+        for i, layer in enumerate(results):
+            for j, sub_layer in enumerate(layer[0]):
+                for k, weight in enumerate(sub_layer):
+                    results[i][0][j][k] = hardthreshold(weight, num_keep)
 
         weighted_weights1 = [
             [layer * num_examples for layer in weights]
diff --git a/baselines/fedht/fedht/utils.py b/baselines/fedht/fedht/utils.py
index ad67623b5d9..54232ccdd0e 100644
--- a/baselines/fedht/fedht/utils.py
+++ b/baselines/fedht/fedht/utils.py
@@ -26,24 +26,24 @@ def partition_data(data, num_partitions):
     X, y = data
     partition_size = len(X) // num_partitions
     # Create partitions
-    partitionsX = [
+    partitions_x = [
         X[i * partition_size : (i + 1) * partition_size] for i in range(num_partitions)
     ]
-    partitionsy = [
+    partitions_y = [
         y[i * partition_size : (i + 1) * partition_size] for i in range(num_partitions)
     ]
 
     # Handle any remaining items
     if len(data) % num_partitions != 0:
         # partitions[-1] = partitions[-1] + data[num_partitions * partition_size:]
-        partitionsX[-1] = np.vstack(
-            (partitionsX[-1], X[num_partitions * partition_size :])
+        partitions_x[-1] = np.vstack(
+            (partitions_x[-1], X[num_partitions * partition_size :])
         )
-        partitionsy[-1] = np.vstack(
-            (partitionsy[-1], y[num_partitions * partition_size :])
+        partitions_y[-1] = np.vstack(
+            (partitions_y[-1], y[num_partitions * partition_size :])
         )
 
-    return partitionsX, partitionsy
+    return partitions_x, partitions_y
 
 
 def sim_data(ni: int, num_clients: int, num_features: int, alpha=1, beta=1):