tasks.py

mkang315 · Mar 25, 2024 · 73048ff · 73048ff
1 parent dcb3e2f
commit 73048ff
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nn/tasks.py b/nn/tasks.py
@@ -10,7 +10,7 @@
 from ...nn.modules import (AIFI, C1, C2, C3, C3TR, SPP, SPPF, Bottleneck, BottleneckCSP, C2f, C3Ghost, C3x,
                                     Classify, Concat, Conv, Conv2, ConvTranspose, Detect, DWConv, DWConvTranspose2d,
                                     Focus, GhostBottleneck, GhostConv, HGBlock, HGStem, RepC3, RepConv,
-                                    cbam_block, eca_block, CA_Block, se_block,CSPStage,BiLevelRoutingAttention)
+                                    RTDETRDecoder, Segment, cbam_block, eca_block, CA_Block, se_block,CSPStage,BiLevelRoutingAttention)
 from ...yolo.utils import DEFAULT_CFG_DICT, DEFAULT_CFG_KEYS, LOGGER, colorstr, emojis, yaml_load
 from ...yolo.utils.checks import check_requirements, check_suffix, check_yaml
 from ...yolo.utils.loss import v8DetectionLoss