I think there might be an issue because the same data tensor is used multiple times, and I am wondering if there is a gradient calculation here. Can you see what happens when you set data.requires_grad = False when it is returned by batch.get?
A miscellaneous note: You can try timing to see if there is an actual difference, but I don’t think the line op = torch.zeros(size) #zero_initializations will actually save time based on how allocations are done. If anything it might just make autograd more confusing if the assignment happens with the actual tensor values rather than just the object references.