import torch
import torch.nn.functional as F

device = torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
print(f"Device: {device}, torch version: {torch.__version__}")

with torch.device(device):
    query = torch.randn(2, 4, 16, 8)
    key = torch.randn(2, 4, 16, 8)
    value = torch.randn(2, 4, 16, 8)

    hidden_states = F.scaled_dot_product_attention(
        query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
    )
    assert torch.isfinite(hidden_states).all()

    query += torch.nan
    hidden_states = F.scaled_dot_product_attention(
        query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False
    )
    assert torch.isnan(hidden_states).all()