from __future__ import annotations

import pytest
import torch

from hippotorch.encoders.dual import DualEncoder


@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_dual_encoder_forward_on_cuda() -> None:
    dev = torch.device("cuda")
    input_dim, embed_dim = 8, 16
    enc = DualEncoder(input_dim=input_dim, embed_dim=embed_dim).to(dev)
    x = torch.randn(1, 2, input_dim, device=dev)
    q = enc.encode_query(x)
    with torch.no_grad():
        k = enc.encode_key(x)
    assert q.shape == (1, embed_dim)
    assert k.shape == (1, embed_dim)
