in src/helpers.py [0:0]
def setup(args: dict, seed: int = 0) -> dict:
args["device"] = "cuda" if torch.cuda.is_available() else "cpu"
# Set random seed
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
return args