diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml
index 112741e3..8b176add 100644
--- a/.github/workflows/tests.yaml
+++ b/.github/workflows/tests.yaml
@@ -14,7 +14,6 @@ jobs:
       fail-fast: false
       matrix:
         os: [ubuntu-latest, windows-latest]
-        version: ['3.10', '3.11']
 
     runs-on: ${{ matrix.os }}
 
@@ -27,7 +26,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
-        python-version: ${{ matrix.version }}
+        python-version: '3.10'
         cache: pip
 
     - name: Install dependencies [Ubuntu]
diff --git a/luxonis_train/nodes/__init__.py b/luxonis_train/nodes/__init__.py
index 4c90abaa..5b4a889f 100644
--- a/luxonis_train/nodes/__init__.py
+++ b/luxonis_train/nodes/__init__.py
@@ -1,37 +1,4 @@
-from .base_node import BaseNode
-from .bisenet_head import BiSeNetHead
-from .classification_head import ClassificationHead
-from .contextspatial import ContextSpatial
-from .efficient_bbox_head import EfficientBBoxHead
-from .efficient_keypoint_bbox_head import EfficientKeypointBBoxHead
-from .efficientnet import EfficientNet
-from .efficientrep import EfficientRep
-from .implicit_keypoint_bbox_head import ImplicitKeypointBBoxHead
-from .micronet import MicroNet
-from .mobilenetv2 import MobileNetV2
-from .mobileone import MobileOne
-from .reppan_neck import RepPANNeck
-from .repvgg import RepVGG
-from .resnet import ResNet
-from .rexnetv1 import ReXNetV1_lite
-from .segmentation_head import SegmentationHead
-
-__all__ = [
-    "BiSeNetHead",
-    "ClassificationHead",
-    "ContextSpatial",
-    "EfficientBBoxHead",
-    "EfficientNet",
-    "EfficientRep",
-    "EfficientKeypointBBoxHead",
-    "ImplicitKeypointBBoxHead",
-    "BaseNode",
-    "MicroNet",
-    "MobileNetV2",
-    "MobileOne",
-    "ReXNetV1_lite",
-    "RepPANNeck",
-    "RepVGG",
-    "ResNet",
-    "SegmentationHead",
-]
+from .backbones import *
+from .base_node import *
+from .heads import *
+from .necks import *
diff --git a/luxonis_train/nodes/backbones/__init__.py b/luxonis_train/nodes/backbones/__init__.py
new file mode 100644
index 00000000..9463124b
--- /dev/null
+++ b/luxonis_train/nodes/backbones/__init__.py
@@ -0,0 +1,21 @@
+from .contextspatial import ContextSpatial
+from .efficientnet import EfficientNet
+from .efficientrep import EfficientRep
+from .micronet import MicroNet
+from .mobilenetv2 import MobileNetV2
+from .mobileone import MobileOne
+from .repvgg import RepVGG
+from .resnet import ResNet
+from .rexnetv1 import ReXNetV1_lite
+
+__all__ = [
+    "ContextSpatial",
+    "EfficientNet",
+    "EfficientRep",
+    "MicroNet",
+    "MobileNetV2",
+    "MobileOne",
+    "ReXNetV1_lite",
+    "RepVGG",
+    "ResNet",
+]
diff --git a/luxonis_train/nodes/contextspatial.py b/luxonis_train/nodes/backbones/contextspatial.py
similarity index 98%
rename from luxonis_train/nodes/contextspatial.py
rename to luxonis_train/nodes/backbones/contextspatial.py
index 1ca1460d..2cac4b81 100644
--- a/luxonis_train/nodes/contextspatial.py
+++ b/luxonis_train/nodes/backbones/contextspatial.py
@@ -7,6 +7,7 @@
 from torch import Tensor, nn
 from torch.nn import functional as F
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import (
     AttentionRefinmentBlock,
     ConvModule,
@@ -14,8 +15,6 @@
 )
 from luxonis_train.utils.registry import NODES
 
-from .base_node import BaseNode
-
 
 class ContextSpatial(BaseNode[Tensor, list[Tensor]]):
     def __init__(self, context_backbone: str = "MobileNetV2", **kwargs):
diff --git a/luxonis_train/nodes/efficientnet.py b/luxonis_train/nodes/backbones/efficientnet.py
similarity index 96%
rename from luxonis_train/nodes/efficientnet.py
rename to luxonis_train/nodes/backbones/efficientnet.py
index 37f8ced5..e560bc5f 100644
--- a/luxonis_train/nodes/efficientnet.py
+++ b/luxonis_train/nodes/backbones/efficientnet.py
@@ -7,7 +7,7 @@
 import torch
 from torch import Tensor, nn
 
-from .base_node import BaseNode
+from luxonis_train.nodes.base_node import BaseNode
 
 
 class EfficientNet(BaseNode[Tensor, list[Tensor]]):
diff --git a/luxonis_train/nodes/efficientrep.py b/luxonis_train/nodes/backbones/efficientrep.py
similarity index 98%
rename from luxonis_train/nodes/efficientrep.py
rename to luxonis_train/nodes/backbones/efficientrep.py
index 24e43397..be558620 100644
--- a/luxonis_train/nodes/efficientrep.py
+++ b/luxonis_train/nodes/backbones/efficientrep.py
@@ -9,6 +9,7 @@
 
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import (
     BlockRepeater,
     RepVGGBlock,
@@ -16,8 +17,6 @@
 )
 from luxonis_train.utils.general import make_divisible
 
-from .base_node import BaseNode
-
 logger = logging.getLogger(__name__)
 
 
diff --git a/luxonis_train/nodes/micronet.py b/luxonis_train/nodes/backbones/micronet.py
similarity index 99%
rename from luxonis_train/nodes/micronet.py
rename to luxonis_train/nodes/backbones/micronet.py
index 603eabde..074dce2a 100644
--- a/luxonis_train/nodes/micronet.py
+++ b/luxonis_train/nodes/backbones/micronet.py
@@ -4,10 +4,9 @@
 from torch import Tensor, nn
 
 from luxonis_train.nodes.activations import HSigmoid, HSwish
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
 
-from .base_node import BaseNode
-
 
 class MicroNet(BaseNode[Tensor, list[Tensor]]):
     """
diff --git a/luxonis_train/nodes/mobilenetv2.py b/luxonis_train/nodes/backbones/mobilenetv2.py
similarity index 95%
rename from luxonis_train/nodes/mobilenetv2.py
rename to luxonis_train/nodes/backbones/mobilenetv2.py
index db6cf879..48161835 100644
--- a/luxonis_train/nodes/mobilenetv2.py
+++ b/luxonis_train/nodes/backbones/mobilenetv2.py
@@ -6,7 +6,7 @@
 import torchvision
 from torch import Tensor, nn
 
-from .base_node import BaseNode
+from luxonis_train.nodes.base_node import BaseNode
 
 
 class MobileNetV2(BaseNode[Tensor, list[Tensor]]):
diff --git a/luxonis_train/nodes/mobileone.py b/luxonis_train/nodes/backbones/mobileone.py
similarity index 99%
rename from luxonis_train/nodes/mobileone.py
rename to luxonis_train/nodes/backbones/mobileone.py
index 645534e4..2d460fd0 100644
--- a/luxonis_train/nodes/mobileone.py
+++ b/luxonis_train/nodes/backbones/mobileone.py
@@ -10,10 +10,9 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule, SqueezeExciteBlock
 
-from .base_node import BaseNode
-
 
 class MobileOne(BaseNode[Tensor, list[Tensor]]):
     """Implementation of MobileOne backbone.
diff --git a/luxonis_train/nodes/repvgg.py b/luxonis_train/nodes/backbones/repvgg.py
similarity index 99%
rename from luxonis_train/nodes/repvgg.py
rename to luxonis_train/nodes/backbones/repvgg.py
index f488a68c..c536c78e 100644
--- a/luxonis_train/nodes/repvgg.py
+++ b/luxonis_train/nodes/backbones/repvgg.py
@@ -6,7 +6,7 @@
 
 from luxonis_train.nodes.blocks import RepVGGBlock
 
-from .base_node import BaseNode
+from ..base_node import BaseNode
 
 logger = logging.getLogger(__name__)
 
diff --git a/luxonis_train/nodes/resnet.py b/luxonis_train/nodes/backbones/resnet.py
similarity index 98%
rename from luxonis_train/nodes/resnet.py
rename to luxonis_train/nodes/backbones/resnet.py
index 3f810100..e4228410 100644
--- a/luxonis_train/nodes/resnet.py
+++ b/luxonis_train/nodes/backbones/resnet.py
@@ -8,7 +8,7 @@
 import torchvision
 from torch import Tensor, nn
 
-from .base_node import BaseNode
+from ..base_node import BaseNode
 
 
 class ResNet(BaseNode[Tensor, list[Tensor]]):
diff --git a/luxonis_train/nodes/rexnetv1.py b/luxonis_train/nodes/backbones/rexnetv1.py
similarity index 99%
rename from luxonis_train/nodes/rexnetv1.py
rename to luxonis_train/nodes/backbones/rexnetv1.py
index 181ad325..6d23857e 100644
--- a/luxonis_train/nodes/rexnetv1.py
+++ b/luxonis_train/nodes/backbones/rexnetv1.py
@@ -7,11 +7,10 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
 from luxonis_train.utils.general import make_divisible
 
-from .base_node import BaseNode
-
 
 class ReXNetV1_lite(BaseNode[Tensor, list[Tensor]]):
     def __init__(
diff --git a/luxonis_train/nodes/base_node.py b/luxonis_train/nodes/base_node.py
index 3b549e0a..9db45316 100644
--- a/luxonis_train/nodes/base_node.py
+++ b/luxonis_train/nodes/base_node.py
@@ -19,6 +19,8 @@
 ForwardOutputT = TypeVar("ForwardOutputT")
 ForwardInputT = TypeVar("ForwardInputT")
 
+__all__ = ["BaseNode"]
+
 
 class BaseNode(
     nn.Module,
diff --git a/luxonis_train/nodes/heads/__init__.py b/luxonis_train/nodes/heads/__init__.py
new file mode 100644
index 00000000..28b5e8ca
--- /dev/null
+++ b/luxonis_train/nodes/heads/__init__.py
@@ -0,0 +1,15 @@
+from .bisenet_head import BiSeNetHead
+from .classification_head import ClassificationHead
+from .efficient_bbox_head import EfficientBBoxHead
+from .efficient_keypoint_bbox_head import EfficientKeypointBBoxHead
+from .implicit_keypoint_bbox_head import ImplicitKeypointBBoxHead
+from .segmentation_head import SegmentationHead
+
+__all__ = [
+    "BiSeNetHead",
+    "ClassificationHead",
+    "EfficientBBoxHead",
+    "EfficientKeypointBBoxHead",
+    "ImplicitKeypointBBoxHead",
+    "SegmentationHead",
+]
diff --git a/luxonis_train/nodes/bisenet_head.py b/luxonis_train/nodes/heads/bisenet_head.py
similarity index 96%
rename from luxonis_train/nodes/bisenet_head.py
rename to luxonis_train/nodes/heads/bisenet_head.py
index 54128cad..3fef7584 100644
--- a/luxonis_train/nodes/bisenet_head.py
+++ b/luxonis_train/nodes/heads/bisenet_head.py
@@ -7,12 +7,11 @@
 
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import ConvModule
 from luxonis_train.utils.general import infer_upscale_factor
 from luxonis_train.utils.types import LabelType, Packet
 
-from .base_node import BaseNode
-
 
 class BiSeNetHead(BaseNode[Tensor, Tensor]):
     in_height: int
diff --git a/luxonis_train/nodes/classification_head.py b/luxonis_train/nodes/heads/classification_head.py
similarity index 94%
rename from luxonis_train/nodes/classification_head.py
rename to luxonis_train/nodes/heads/classification_head.py
index d33faeb5..07b3d72b 100644
--- a/luxonis_train/nodes/classification_head.py
+++ b/luxonis_train/nodes/heads/classification_head.py
@@ -1,9 +1,8 @@
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.utils.types import LabelType
 
-from .base_node import BaseNode
-
 
 class ClassificationHead(BaseNode[Tensor, Tensor]):
     in_channels: int
diff --git a/luxonis_train/nodes/efficient_bbox_head.py b/luxonis_train/nodes/heads/efficient_bbox_head.py
similarity index 99%
rename from luxonis_train/nodes/efficient_bbox_head.py
rename to luxonis_train/nodes/heads/efficient_bbox_head.py
index e80ca31a..5607a2a8 100644
--- a/luxonis_train/nodes/efficient_bbox_head.py
+++ b/luxonis_train/nodes/heads/efficient_bbox_head.py
@@ -9,6 +9,7 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import EfficientDecoupledBlock
 from luxonis_train.utils.boxutils import (
     anchors_for_fpn_features,
@@ -17,8 +18,6 @@
 )
 from luxonis_train.utils.types import LabelType, Packet
 
-from .base_node import BaseNode
-
 
 class EfficientBBoxHead(
     BaseNode[list[Tensor], tuple[list[Tensor], list[Tensor], list[Tensor]]]
diff --git a/luxonis_train/nodes/efficient_keypoint_bbox_head.py b/luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
similarity index 100%
rename from luxonis_train/nodes/efficient_keypoint_bbox_head.py
rename to luxonis_train/nodes/heads/efficient_keypoint_bbox_head.py
diff --git a/luxonis_train/nodes/implicit_keypoint_bbox_head.py b/luxonis_train/nodes/heads/implicit_keypoint_bbox_head.py
similarity index 99%
rename from luxonis_train/nodes/implicit_keypoint_bbox_head.py
rename to luxonis_train/nodes/heads/implicit_keypoint_bbox_head.py
index 79e3fb79..0ca995c5 100644
--- a/luxonis_train/nodes/implicit_keypoint_bbox_head.py
+++ b/luxonis_train/nodes/heads/implicit_keypoint_bbox_head.py
@@ -5,6 +5,7 @@
 import torch
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import KeypointBlock, LearnableMulAddConv
 from luxonis_train.utils.boxutils import (
     non_max_suppression,
@@ -13,8 +14,6 @@
 )
 from luxonis_train.utils.types import LabelType, Packet
 
-from .base_node import BaseNode
-
 logger = logging.getLogger(__name__)
 
 
diff --git a/luxonis_train/nodes/segmentation_head.py b/luxonis_train/nodes/heads/segmentation_head.py
similarity index 96%
rename from luxonis_train/nodes/segmentation_head.py
rename to luxonis_train/nodes/heads/segmentation_head.py
index cddfe8db..1b29df7b 100644
--- a/luxonis_train/nodes/segmentation_head.py
+++ b/luxonis_train/nodes/heads/segmentation_head.py
@@ -7,12 +7,11 @@
 import torch.nn as nn
 from torch import Tensor
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import UpBlock
 from luxonis_train.utils.general import infer_upscale_factor
 from luxonis_train.utils.types import LabelType
 
-from .base_node import BaseNode
-
 
 class SegmentationHead(BaseNode[Tensor, Tensor]):
     in_height: int
diff --git a/luxonis_train/nodes/necks/__init__.py b/luxonis_train/nodes/necks/__init__.py
new file mode 100644
index 00000000..eef2e9a0
--- /dev/null
+++ b/luxonis_train/nodes/necks/__init__.py
@@ -0,0 +1,3 @@
+from .reppan_neck import RepPANNeck
+
+__all__ = ["RepPANNeck"]
diff --git a/luxonis_train/nodes/reppan_neck.py b/luxonis_train/nodes/necks/reppan_neck.py
similarity index 99%
rename from luxonis_train/nodes/reppan_neck.py
rename to luxonis_train/nodes/necks/reppan_neck.py
index 26fed274..bd05f083 100644
--- a/luxonis_train/nodes/reppan_neck.py
+++ b/luxonis_train/nodes/necks/reppan_neck.py
@@ -10,11 +10,10 @@
 
 from torch import Tensor, nn
 
+from luxonis_train.nodes.base_node import BaseNode
 from luxonis_train.nodes.blocks import RepDownBlock, RepUpBlock
 from luxonis_train.utils.general import make_divisible
 
-from .base_node import BaseNode
-
 
 class RepPANNeck(BaseNode[list[Tensor], list[Tensor]]):
     def __init__(
diff --git a/media/coverage_badge.svg b/media/coverage_badge.svg
index 12876e69..8e21255a 100644
--- a/media/coverage_badge.svg
+++ b/media/coverage_badge.svg
@@ -15,7 +15,7 @@
     <g fill="#fff" text-anchor="middle" font-family="DejaVu Sans,Verdana,Geneva,sans-serif" font-size="11">
         <text x="31.5" y="15" fill="#010101" fill-opacity=".3">coverage</text>
         <text x="31.5" y="14">coverage</text>
-        <text x="80" y="15" fill="#010101" fill-opacity=".3">78%</text>
-        <text x="80" y="14">78%</text>
+        <text x="80" y="15" fill="#010101" fill-opacity=".3">84%</text>
+        <text x="80" y="14">84%</text>
     </g>
 </svg>
diff --git a/tests/integration/test_sanity.py b/tests/integration/test_sanity.py
index 1185bb8d..c6d3bf09 100644
--- a/tests/integration/test_sanity.py
+++ b/tests/integration/test_sanity.py
@@ -35,9 +35,17 @@ def clear_output():
 
 
 @pytest.mark.parametrize(
-    "config_file", [str(path) for path in Path("configs").glob("*model*")]
+    "config_file",
+    [
+        "classification_model",
+        "segmentation_model",
+        "detection_model",
+        "keypoint_bbox_model",
+        "resnet_model",
+    ],
 )
 def test_simple_models(config_file: str):
+    config_file = f"configs/{config_file}.yaml"
     model = LuxonisModel(config_file, opts=OPTS)
     model.train()
     model.test()