def get_gpu_memory(gpus):
    """
    Small heuristic to calculate the amount of memory needed for each GPU in case of multi-gpu training.
    """
    if len(gpus) >= 5:
        return 0.4
    elif len(gpus) >= 3:
        return 0.5
    elif len(gpus) == 2:
        return 0.6
    else:
        return 0.6