diff --git a/luxonis_train/nodes/__init__.py b/luxonis_train/nodes/__init__.py index 954db2be..9a506c1f 100644 --- 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 index 0b0aedde..57b52d09 100644 --- 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.