Why your EmbeddingGemma 2 vectors are NaN (and the one-line fix)

Loading a model with torch_dtype=torch.float16 is a common habit for saving memory. With EmbeddingGemma 2 it’s a trap: the activations go beyond float16’s range. You don’t get an error. You get NaNs, or embeddings that look fine and rank badly.

The fix

import torch
from sentence_transformers import SentenceTransformer

use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
dtype = torch.bfloat16 if use_bf16 else torch.float32
model = SentenceTransformer("google/embeddinggemma-2", model_kwargs={"torch_dtype": dtype})

bfloat16 has the same 8-bit exponent as float32, so it uses half the memory without overflowing. Most CPUs and older GPUs (pre-Ampere, such as T4 and V100) don’t support bf16 natively, so use float32 there.

A quick sanity check

Add this once to your pipeline:

import numpy as np

emb = model.encode(["hello world"])
assert not np.isnan(emb).any(), "NaN embeddings: check your dtype"

Does this apply to other models?

Often, yes. Models in the Gemma family usually ship with bf16 weights. Whenever a model card says “bf16 or fp32, not fp16”, take it at its word.