Style fixes

src/utils/cache.py

@@ -45,7 +45,7 @@ def save(path, file_id, indices) -> None:
     @staticmethod
     def load(path, file_id, device="cpu") -> Tensor:
         file_path = os.path.join(path, file_id)
-        return torch.load(file_path, map_location = device)
+        return torch.load(file_path, map_location=device)
 
     @staticmethod
     def exists(path, file_id) -> bool:

src/utils/datasets/mark_dataset.py

@@ -1,4 +1,4 @@
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Union
 
 import torch
 from torch.utils.data.dataset import Dataset
@@ -16,7 +16,7 @@ def __init__(
         cache_path: str = "./datasets",
         p: float = 0.3,
         cls_to_mark: int = 2,
-        mark_fn: Optional[Union[Callable, str]]=None,
+        mark_fn: Optional[Union[Callable, str]] = None,
         only_train: bool = False
     ):
         super().__init__()
@@ -28,9 +28,9 @@ def __init__(
         self.cls_to_mark = cls_to_mark
         self.mark_prob = p
         if mark_fn is not None:
-            self.mark_image=mark_fn
+            self.mark_image = mark_fn
         else:
-            self.mark_image=self.mark_image_contour_and_square
+            self.mark_image = self.mark_image_contour_and_square
 
         if IC.exists(path=cache_path, file_id=f"{dataset_id}_mark_ids"):
             self.mark_indices = IC.load(path="./datasets", file_id=f"{dataset_id}_mark_ids")
@@ -44,7 +44,7 @@ def __len__(self):
     def __getitem__(self, index):
         x, y = self.dataset[index]
         if index in self.mark_indices:
-            x=self.inverse_transform(x)
+            x = self.inverse_transform(x)
             return self.dataset.transform(self.mark_image(x)), y
         else:
             return x, y