#!/usr/bin/env python3
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
import megengine as mge
import megengine.module as M
import numpy as np
import pytest
import torch
import torch.nn as nn

from basecls.configs import BaseConfig
from basecls.layers import BinaryCrossEntropy, CrossEntropy, build_loss


@pytest.mark.parametrize("name", [CrossEntropy, "BinaryCrossEntropy", "CrossEntropy"])
def test_build_loss(name):
    cfg = BaseConfig(loss=dict(name=name))

    m = build_loss(cfg)
    assert isinstance(m, M.Module)


def test_bce():
    x = np.random.rand(2, 8, 4).astype("float32")
    y = np.random.rand(2, 8, 4).astype("float32")

    ml = BinaryCrossEntropy()(
mge.Tensor(x)