fixed efficientnet

luxonis · Feb 22, 2024 · bd0dfe2 · bd0dfe2
1 parent 79df5c6
commit bd0dfe2
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/luxonis_train/nodes/__init__.py b/luxonis_train/nodes/__init__.py
@@ -3,6 +3,7 @@
 from .classification_head import ClassificationHead
 from .contextspatial import ContextSpatial
 from .efficient_bbox_head import EfficientBBoxHead
+from .efficientnet import EfficientNet
 from .efficientrep import EfficientRep
 from .implicit_keypoint_bbox_head import ImplicitKeypointBBoxHead
 from .micronet import MicroNet
@@ -19,6 +20,7 @@
     "ClassificationHead",
     "ContextSpatial",
     "EfficientBBoxHead",
+    "EfficientNet",
     "EfficientRep",
     "ImplicitKeypointBBoxHead",
     "BaseNode",

diff --git a/luxonis_train/nodes/efficientnet.py b/luxonis_train/nodes/efficientnet.py
@@ -11,6 +11,8 @@
 
 
 class EfficientNet(BaseNode[Tensor, list[Tensor]]):
+    attach_index: int = -1
+
     def __init__(self, download_weights: bool = False, **kwargs):
         """EfficientNet backbone.