Repository: qdrant/quaterion-models (language: python)

Issue:
Fix imports brought in #2



Discussion / hints:

Fix patch:
diff --git a/quaterion_models/encoders/encoder.py b/quaterion_models/encoders/encoder.py
index c55c2cb..1e9853d 100644
--- a/quaterion_models/encoders/encoder.py
+++ b/quaterion_models/encoders/encoder.py
@@ -1,10 +1,9 @@
-from typing import List, Any, Union, Dict, Tuple, Callable
-
 from torch import Tensor
 from torch import nn
 
 from quaterion_models.types import TensorInterchange, CollateFnType
 
+
 class Encoder(nn.Module):
     """
     Base class for encoder abstraction
diff --git a/quaterion_models/heads/__init__.py b/quaterion_models/heads/__init__.py
index b2bba13..f65f20a 100644
--- a/quaterion_models/heads/__init__.py
+++ b/quaterion_models/heads/__init__.py
@@ -1,4 +1,4 @@
-from quaterion_models.heads.empty_head import EmptyHead
 from quaterion_models.heads.encoder_head import EncoderHead
+from quaterion_models.heads.empty_head import EmptyHead
 from quaterion_models.heads.gated_head import GatedHead
 from quaterion_models.heads.softmax_head import SoftmaxEmbeddingsHead
diff --git a/quaterion_models/heads/empty_head.py b/quaterion_models/heads/empty_head.py
index 7a15e9b..457eb1c 100644
--- a/quaterion_models/heads/empty_head.py
+++ b/quaterion_models/heads/empty_head.py
@@ -1,5 +1,3 @@
-from typing import Dict, Any
-
 import torch
 
 from quaterion_models.heads import EncoderHead
diff --git a/quaterion_models/types/__init__.py b/quaterion_models/types/__init__.py
index e15d71c..ee0b197 100644
--- a/quaterion_models/types/__init__.py
+++ b/quaterion_models/types/__init__.py
@@ -1,5 +1,14 @@
 # TODO: Split them into separate files once we have more of them.
+from typing import Union, Tuple, List, Dict, Any, Callable
+
+from torch import Tensor
+
 TensorInterchange = Union[
-    Tensor, Tuple[Tensor], List[Tensor], Dict[str, Tensor], Dict[str, dict], Any
+    Tensor,
+    Tuple[Tensor],
+    List[Tensor],
+    Dict[str, Tensor],
+    Dict[str, dict],
+    Any,
 ]
 CollateFnType = Callable[[List[Any]], TensorInterchange]


Test patch:
diff --git a/tests/test_switch_encoder.py b/tests/test_switch_encoder.py
index b619e16..a203e34 100644
--- a/tests/test_switch_encoder.py
+++ b/tests/test_switch_encoder.py
@@ -5,8 +5,9 @@
 import torch
 from torch import Tensor
 
-from quaterion_models.encoder import Encoder, TensorInterchange, CollateFnType
+from quaterion_models.encoders import Encoder
 from quaterion_models.heads.empty_head import EmptyHead
+from quaterion_models.types import TensorInterchange, CollateFnType
 from quaterion_models.model import MetricModel
 from quaterion_models.encoders.switch_encoder import SwitchEncoder