src/nn/criterion/__init__.py

import torch.nn as nn 
from src.core import register

CrossEntropyLoss = register(nn.CrossEntropyLoss)