Function _enough_gpu_memory

tests/test_embedding.py:169–173

(num_bytes: int)

# logic; this one validates the real index range (incl. boundary ids 0 and
# vocab-1) and the real hidden width, with a small (batch, seq) load point.
def _enough_gpu_memory(num_bytes: int) -> bool:
    if not torch.cuda.is_available():
        return False
    free, _ = torch.cuda.mem_get_info()
    return free > int(num_bytes * 1.5)  # headroom for the gathered output


@pytest.mark.skipif(