import torch
import torch.nn as nn
import math

N, C, H, W = 32, 64, 56, 56


class LogCoshLoss(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.reduction = reduction

    def forward(self, input, target):
        diff = input - target
        loss = torch.abs(diff) + torch.nn.functional.softplus(-2. * torch.abs(diff)) - math.log(2.0)

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss


class Model(nn.Module):
    def __init__(self, reduction='mean'):
        super().__init__()
        self.op = LogCoshLoss(reduction)

    def forward(self, input, target):
        return self.op(input, target)


def get_inputs():
    input = torch.randn(N, C, H, W, dtype=torch.float32)
    target = torch.randn(N, C, H, W, dtype=torch.float32)
    return [input, target]


def get_init_inputs():
    return ['mean']