PyTorch Contrastive Loss

Task

Implement class ContrastiveRetriever(query_weight: torch.Tensor, document_weight: torch.Tensor, learning_rate: float, temperature: float). Debug a small end-to-end PyTorch contrastive retriever. The class should define trainable query/document projection parameters, produce forward logits, compute diagonal contrastive loss, run conventional backward/optimizer steps, and fit over a small sample dataset for several mini-batch epochs.

Requirements

• The starter implementation is intentionally buggy and can raise a matmul shape mismatch before its logic bugs are repaired.
• Subclass torch.nn.Module and call super().__init__() before assigning parameters.
• Clone query_weight and document_weight into torch.nn.Parameter objects named query_projection and document_projection.
• Create torch.optim.SGD from self.parameters() using learning_rate.
• forward(query_features, document_features) projects rows with features @ projection, L2-normalizes along dim=-1, and returns query-document logits divided by temperature.
• loss(logits) uses diagonal integer targets [0, 1, ..., batch_size - 1] on the logits device and returns mean cross-entropy.
• backward(loss) backpropagates one scalar loss without detaching it.
• training_step(query_batch, document_batch) runs optimizer.zero_grad(), forward, loss, backward(loss), optimizer.step(), and returns the detached pre-update loss.
• fit(query_features, document_features, batch_size, epochs) loops through every sequential mini-batch for every epoch, appends each pre-update loss to loss_history, and returns loss_history as a tensor.
• Do not mutate the caller's initial weight tensors in-place.

Example

model = ContrastiveRetriever(torch.eye(2), torch.eye(2), 0.1, 1.0)
model.fit(torch.eye(2), torch.eye(2), batch_size=2, epochs=1)
# tensor([0.3133])