import torch
from torch import FloatTensor, LongTensor
from torch.autograd import Variable
from torch import nn, optim
from torch.nn import Parameter
from tqdm import trange

import util, logging, os, psutil

import hyper, gaussian

logging.basicConfig(filename='memleak.log',level=logging.INFO)

torch.manual_seed(2)

CUDA = True

B = 256
M = 32
W, H = 2, 2048

sparsemult = util.SparseMult(use_cuda=CUDA)

def iteration():

    mindices = (torch.rand(B, H, W) * M).long()
    values = torch.rand(B, H)
    bsize = LongTensor([M, M])
    y_flat = torch.zeros(B, M)

    x_flat = torch.rand(B, M)

    if CUDA:
        mindices = mindices.cuda()
        bsize = bsize.cuda()
        values = values.cuda()
        x_flat = x_flat.cuda()
        y_flat = y_flat.cuda()

    bsize = Variable(bsize)
    values = Variable(values)
    x_flat = Variable(x_flat)
    y_flat = Variable(y_flat)

    if CUDA:
        logging.info(util.nvidia_smi())

    for b in range(B):
        bindices = Variable(mindices[b, :, :].squeeze(0).t())
        bvalues = values[b, :]

        bx = x_flat[b, :]

        y_flat[b, :] = sparsemult(bindices, bvalues, bsize, bx)

for i in trange(int(10e7)):
    iteration()

@@ -399,16 +399,41 @@
 da=CUDA)
+%0Acriterion = nn.MSELoss()
 %0A%0Adef it
@@ -483,16 +483,18 @@
 , W) * M
+-1
 ).long()
@@ -582,25 +582,24 @@
 zeros(B, M)%0A
-%0A
     x_flat =
@@ -615,16 +615,44 @@
 nd(B, M)
+%0A    target = x_flat.clone()
 %0A%0A    if
@@ -814,16 +814,47 @@
 t.cuda()
+%0A        target = target.cuda()
 %0A%0A    bs
@@ -903,16 +903,36 @@
 e(values
+, requires_grad=True
 )%0A    x_
@@ -984,16 +984,46 @@
 (y_flat)
+%0A    target = Variable(target)
 %0A%0A    if
@@ -1279,16 +1279,91 @@
 e, bx)%0A%0A
+    loss = criterion(y_flat, target)%0A    loss.backward(retain_graph=True)%0A%0A
 for i in

Add backward to repro script