@@ -1,5 +1,5 @@
 from super_gradients.common.factories.base_factory import BaseFactory
-from super_gradients.training.losses import LOSSES
+from super_gradients.common.registry.registry import LOSSES
 
 
 class LossesFactory(BaseFactory):

@@ -1,12 +1,14 @@
 import inspect
 from typing import Callable, Dict, Optional
 
+from torch import nn
+
+from super_gradients.common import object_names
 from super_gradients.training.utils.callbacks import LR_SCHEDULERS_CLS_DICT
 from super_gradients.common.sg_loggers import SG_LOGGERS
 from super_gradients.training.dataloaders.dataloaders import ALL_DATALOADERS
 from super_gradients.training.models.all_architectures import ARCHITECTURES
 from super_gradients.training.metrics.all_metrics import METRICS
-from super_gradients.training.losses.all_losses import LOSSES
 from super_gradients.modules.detection_modules import ALL_DETECTION_MODULES
 from super_gradients.training.utils.callbacks.all_callbacks import CALLBACKS
 from super_gradients.training.transforms.all_transforms import TRANSFORMS
@@ -24,8 +26,8 @@ def create_register_decorator(registry: Dict[str, Callable]) -> Callable:
     """
     Create a decorator that registers object of specified type (model, metric, ...)
 
-    :param registry: The registry (maps name to object that you register)
-    :return:         Register function
+    :param registry:    Dict including registered objects (maps name to object that you register)
+    :return:            Register function
     """
 
     def register(name: Optional[str] = None) -> Callable:
@@ -55,7 +57,10 @@ def create_register_decorator(registry: Dict[str, Callable]) -> Callable:
 register_model = create_register_decorator(registry=ARCHITECTURES)
 register_detection_module = create_register_decorator(registry=ALL_DETECTION_MODULES)
 register_metric = create_register_decorator(registry=METRICS)
+
+LOSSES = {object_names.Losses.MSE: nn.MSELoss}
 register_loss = create_register_decorator(registry=LOSSES)
+
 register_dataloader = create_register_decorator(registry=ALL_DATALOADERS)
 register_callback = create_register_decorator(registry=CALLBACKS)
 register_transform = create_register_decorator(registry=TRANSFORMS)

@@ -1,3 +1,6 @@
+from super_gradients.common.registry.registry import LOSSES
+from super_gradients.common.object_names import Losses
+
 from super_gradients.training.losses.focal_loss import FocalLoss
 from super_gradients.training.losses.kd_losses import KDLogitsLoss
 from super_gradients.training.losses.label_smoothing_cross_entropy_loss import LabelSmoothingCrossEntropyLoss
@@ -8,7 +11,6 @@ from super_gradients.training.losses.yolox_loss import YoloXDetectionLoss, YoloX
 from super_gradients.training.losses.ssd_loss import SSDLoss
 from super_gradients.training.losses.bce_dice_loss import BCEDiceLoss
 from super_gradients.training.losses.dice_ce_edge_loss import DiceCEEdgeLoss
-from super_gradients.training.losses.all_losses import LOSSES, Losses
 from super_gradients.training.losses.ppyolo_loss import PPYoloELoss
 
 __all__ = [

from torch import nn
from super_gradients.common.object_names import Losses
from super_gradients.training.losses import (
    LabelSmoothingCrossEntropyLoss,
    ShelfNetOHEMLoss,
    ShelfNetSemanticEncodingLoss,
    RSquaredLoss,
    SSDLoss,
    BCEDiceLoss,
    YoloXDetectionLoss,
    YoloXFastDetectionLoss,
    KDLogitsLoss,
    DiceCEEdgeLoss,
)
from super_gradients.training.losses.stdc_loss import STDCLoss
from super_gradients.training.losses.ppyolo_loss import PPYoloELoss
from super_gradients.training.losses.dekr_loss import DEKRLoss


LOSSES = {
    Losses.CROSS_ENTROPY: LabelSmoothingCrossEntropyLoss,
    Losses.MSE: nn.MSELoss,
    Losses.R_SQUARED_LOSS: RSquaredLoss,
    Losses.SHELFNET_OHEM_LOSS: ShelfNetOHEMLoss,
    Losses.SHELFNET_SE_LOSS: ShelfNetSemanticEncodingLoss,
    Losses.YOLOX_LOSS: YoloXDetectionLoss,
    Losses.YOLOX_FAST_LOSS: YoloXFastDetectionLoss,
    Losses.SSD_LOSS: SSDLoss,
    Losses.STDC_LOSS: STDCLoss,
    Losses.BCE_DICE_LOSS: BCEDiceLoss,
    Losses.KD_LOSS: KDLogitsLoss,
    Losses.DICE_CE_EDGE_LOSS: DiceCEEdgeLoss,
    Losses.PPYOLOE_LOSS: PPYoloELoss,
    Losses.DEKR_LOSS: DEKRLoss,
}

@@ -1,9 +1,12 @@
 import torch
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
 from super_gradients.training.losses.bce_loss import BCE
 from super_gradients.training.losses.dice_loss import BinaryDiceLoss
 
 
+@register_loss(Losses.BCE_DICE_LOSS)
 class BCEDiceLoss(torch.nn.Module):
     """
     Binary Cross Entropy + Dice Loss

@@ -3,7 +3,11 @@ from typing import Tuple
 import torch
 from torch import Tensor, nn
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
 
+
+@register_loss(Losses.DEKR_LOSS)
 class DEKRLoss(nn.Module):
     """
     Implementation of the loss function from the "Bottom-Up Human Pose Estimation Via Disentangled Keypoint Regression"

@@ -4,9 +4,13 @@ from super_gradients.training.losses.dice_loss import DiceLoss, BinaryDiceLoss
 from super_gradients.training.utils.segmentation_utils import target_to_binary_edge
 from torch.nn.modules.loss import _Loss
 from typing import Union, Tuple
+
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
 from super_gradients.training.losses.mask_loss import MaskAttentionLoss
 
 
+@register_loss(Losses.DICE_CE_EDGE_LOSS)
 class DiceCEEdgeLoss(_Loss):
     def __init__(
         self,

@@ -1,6 +1,9 @@
 from torch.nn.modules.loss import _Loss, KLDivLoss
 import torch
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
+
 
 class KDklDivLoss(KLDivLoss):
     """KL divergence wrapper for knowledge distillation"""
@@ -12,6 +15,7 @@ class KDklDivLoss(KLDivLoss):
         return super(KDklDivLoss, self).forward(torch.log_softmax(student_output, dim=1), torch.softmax(teacher_output, dim=1))
 
 
+@register_loss(Losses.KD_LOSS)
 class KDLogitsLoss(_Loss):
     """Knowledge distillation loss, wraps the task loss and distillation loss"""
 

@@ -2,6 +2,9 @@ import torch
 from torch import nn
 import torch.nn.functional as F
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
+
 
 def onehot(indexes, N=None, ignore_index=None):
     """
@@ -80,6 +83,7 @@ def cross_entropy(inputs, target, weight=None, ignore_index=-100, reduction="mea
     return loss
 
 
+@register_loss(Losses.CROSS_ENTROPY)
 class LabelSmoothingCrossEntropyLoss(nn.CrossEntropyLoss):
     """CrossEntropyLoss - with ability to recieve distrbution as targets, and optional label smoothing"""
 

@@ -6,6 +6,8 @@ import torch.nn.functional as F
 from torch import nn, Tensor
 
 import super_gradients
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
 from super_gradients.training.datasets.data_formats.bbox_formats.cxcywh import cxcywh_to_xyxy
 from super_gradients.training.utils.bbox_utils import batch_distance2bbox
 from super_gradients.training.utils.distributed_training_utils import (
@@ -640,6 +642,7 @@ class GIoULoss(object):
         return loss * self.loss_weight
 
 
+@register_loss(Losses.PPYOLOE_LOSS)
 class PPYoloELoss(nn.Module):
     def __init__(
         self,

@@ -7,6 +7,11 @@ from torch.nn.modules.loss import _Loss
 from super_gradients.training.utils import convert_to_tensor
 
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
+
+
+@register_loss(Losses.R_SQUARED_LOSS)
 class RSquaredLoss(_Loss):
     def forward(self, output, target):
         # FIXME - THIS NEEDS TO BE CHANGED SUCH THAT THIS CLASS INHERETS FROM _Loss (TAKE A LOOK AT YoLoV3DetectionLoss)

@@ -1,8 +1,11 @@
 import torch
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
 from super_gradients.training.losses.ohem_ce_loss import OhemCELoss
 
 
+@register_loss(Losses.SHELFNET_OHEM_LOSS)
 class ShelfNetOHEMLoss(OhemCELoss):
     def __init__(self, threshold: float = 0.7, mining_percent: float = 1e-4, ignore_lb: int = 255):
         """

@@ -3,6 +3,11 @@ from torch import nn
 from torch.autograd import Variable
 
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
+
+
+@register_loss(Losses.SHELFNET_SE_LOSS)
 class ShelfNetSemanticEncodingLoss(nn.CrossEntropyLoss):
     """2D Cross Entropy Loss with Auxilary Loss"""
 

@@ -4,6 +4,8 @@ import torch
 from torch import nn
 from torch.nn.modules.loss import _Loss
 
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
 from super_gradients.training.utils.detection_utils import calculate_bbox_iou_matrix
 from super_gradients.training.utils.ssd_utils import DefaultBoxes
 
@@ -50,6 +52,7 @@ class HardMiningCrossEntropyLoss(_Loss):
         return closs
 
 
+@register_loss(Losses.SSD_LOSS)
 class SSDLoss(_Loss):
     """
         Implements the loss as the sum of the followings:

@@ -1,11 +1,15 @@
+from typing import Union, Tuple
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from super_gradients.training.utils.segmentation_utils import to_one_hot
 from torch.nn.modules.loss import _Loss
+
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
+from super_gradients.training.utils.segmentation_utils import to_one_hot
 from super_gradients.training.losses.ohem_ce_loss import OhemCELoss, OhemBCELoss, OhemLoss
 from super_gradients.training.losses.dice_loss import BinaryDiceLoss
-from typing import Union, Tuple
 
 
 class DetailAggregateModule(nn.Module):
@@ -107,6 +111,7 @@ class DetailLoss(_Loss):
         return self.weights[0] * bce_loss + self.weights[1] * dice_loss
 
 
+@register_loss(Losses.STDC_LOSS)
 class STDCLoss(_Loss):
     """
     Loss class of STDC-Seg training.

@@ -12,6 +12,9 @@ from torch import nn
 from torch.nn.modules.loss import _Loss
 import torch.nn.functional as F
 
+
+from super_gradients.common.object_names import Losses
+from super_gradients.common.registry.registry import register_loss
 from super_gradients.common.abstractions.abstract_logger import get_logger
 from super_gradients.training.utils import torch_version_is_greater_or_equal
 from super_gradients.training.utils.detection_utils import calculate_bbox_iou_matrix
@@ -79,6 +82,7 @@ class IOUloss(nn.Module):
         return loss
 
 
+@register_loss(Losses.YOLOX_LOSS)
 class YoloXDetectionLoss(_Loss):
     """
     Calculate YOLOX loss:
@@ -607,6 +611,7 @@ class YoloXDetectionLoss(_Loss):
         return num_fg, gt_matched_classes, pred_ious_this_matching, matched_gt_inds
 
 
+@register_loss(Losses.YOLOX_FAST_LOSS)
 class YoloXFastDetectionLoss(YoloXDetectionLoss):
     """
     A completely new implementation of YOLOX loss.