import torch
import torch.nn as nn

class AutomaticWeightedLoss(nn.Module):
    """
    Automatically weighted multi-task loss module.

    This module automatically weights multiple loss functions based on their individual magnitudes.

    Params:
        num (int): The number of loss functions.
    
    """
    def __init__(self, num=2):
        super(AutomaticWeightedLoss, self).__init__()
        params = torch.ones(num, requires_grad=True)
        self.params = torch.nn.Parameter(params)

    def forward(self, *x):
        """
        Forward pass through the AutomaticWeightedLoss module.

        Parameters:
            *x (Tensor): Variable number of loss functions.

        Returns:
            Tensor: Weighted sum of the loss functions.
        """
        loss_sum = 0
        for i, loss in enumerate(x):
            loss_sum += 0.5 / (self.params[i] ** 2) * loss + torch.log(1 + self.params[i] ** 2)
        return loss_sum

if __name__ == '__main__':
    awl = AutomaticWeightedLoss(2)
    params = list(next(awl.parameters()))
