Hello,
I am trying to calculate loss on a function which looks like this:

W is a vector of shape (314287, 300), whereas V is a vector of shape (314287, 3000)… (I am trying to learn a sparse representation).
The relevant portion of my code is as follows:
model = torch.optim.Adam([sparse_embeds], lr=1e-4)
embeds = embeds.cuda()
sparse_embeds = sparse_embeds.cuda()
buckets_size = 1000

def bucketize(tensor, buckets):
  total_buckets = int(tensor.shape[0] / buckets) + 1
  i = 0
  while i*buckets < tensor.shape[0]:
    yield i*buckets, (i+1)*buckets
    i += 1


for e in range(50):
  for i in range(len(embeds)):
    for se in bucketize(embeds, buckets_size):
      start, end = se
      all_dots = embeds[start:end] @ embeds[i]
      s1 = (sparse_embeds[start:end] @ sparse_embeds[i]) - all_dots
      loss = torch.sum((s1.pow(2)))

      print(loss.item())
      model.zero_grad()
      loss.backward()
      model.step()   # <<<<<<<<< Memory crashes here.

I had thought that breaking the loss calculation in multiple buckets would have solved the issue but it doesn’t seem to. Any help at all would be appreciated.