That the grad_fn is different between the first loss and the rest is a hint that something is unexpected. Can you check the shapes of embeddings_batch and labels_batch (maybe check few values too) before they are passed to the loss function?
However, a more fundamental issue is that the current problem setup will still leave all of the intermediate tensors from the forward pass during the backward pass in memory. Unless your loss function has some esoteric scaling properties (e.g., nonlinear memory increase with batch size), I’m not sure this will actually save GPU memory. There are fancy methods to save GPU memory at the cost of a little extra recomputation (e.g., [2006.09616] Dynamic Tensor Rematerialization 1) but I don’t think that is something that is easy to do in upstream PyTorch out of the box.