iyedjb/self-forcing

1import numpy as np2import random3import torch4 5 6def set_seed(seed: int, deterministic: bool = False):7    """8    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.9 10    Args:11        seed (`int`):12            The seed to set.13        deterministic (`bool`, *optional*, defaults to `False`):14            Whether to use deterministic algorithms where available. Can slow down training.15    """16    random.seed(seed)17    np.random.seed(seed)18    torch.manual_seed(seed)19    torch.cuda.manual_seed_all(seed)20 21    if deterministic:22        torch.use_deterministic_algorithms(True)23 24 25def merge_dict_list(dict_list):26    if len(dict_list) == 1:27        return dict_list[0]28 29    merged_dict = {}30    for k, v in dict_list[0].items():31        if isinstance(v, torch.Tensor):32            if v.ndim == 0:33                merged_dict[k] = torch.stack([d[k] for d in dict_list], dim=0)34            else:35                merged_dict[k] = torch.cat([d[k] for d in dict_list], dim=0)36        else:37            # for non-tensor values, we just copy the value from the first item38            merged_dict[k] = v39    return merged_dict40