Handwritten InfoNCE in PyTorch

Requirements Implement the following PyTorch function: The inputs are query embeddings q with shape (B, D) and key embeddings k with shape (B, D). For each index i, (q_i, k_i) is the positive pair, while the remaining B - 1 keys serve as negatives. Normalize both the query embeddings and the key embeddings before calculating their dot-product logits. Use temperature tau to compute the symmetric InfoNCE objective, averaging the query-to-key and key-to-query losses. Return one…