from typing import Union
import torch


def convert_dict_entries_to_tensors(
    loss_params: dict, device: Union[str, torch.device]
) -> dict:
    """Set scalars in loss to torch tensors for use with unsupervised losses.

    Args:
        loss_params: loss dictionary to loop over
        device: device to send floats and ints to

    Returns:
        dict with updated values

    """
    for loss, params in loss_params.items():
        for key, val in params.items():
            if type(val) == float:
                loss_params[loss][key] = torch.tensor(
                    val, dtype=torch.float, device=device
                )
            if type(val) == int:
                loss_params[loss][key] = torch.tensor(
                    val, dtype=torch.int, device=device
                )
    return loss_params