ruff fix

leakpro/fl_utils/data_utils.py

@@ -2,6 +2,7 @@
 from torch import Tensor, cat, mean, randn, std, tensor
 from torch.utils.data import DataLoader, Dataset, TensorDataset
 
+
 def get_meanstd(trainset: Dataset, axis_to_reduce: tuple=(-2,-1)) -> tuple[Tensor, Tensor]:
     """Get mean and std of a dataset."""
     cc = cat([trainset[i][0].unsqueeze(0) for i in range(len(trainset))], dim=0)

leakpro/run.py

@@ -1,7 +1,5 @@
 """Run script."""
-import os
-from collections import OrderedDict
-from typing import Callable, Optional
+from typing import Callable
 
 import torch
 from torch import Tensor