@dpernes I am facing a similar issue, but I am not quite able to understand the logic why it is failing? I have implemented a custom regularizer, which calculates the infinity norm of the weights of the model and then adds it to the existing loss.
def l2_rego(mdl):
        l2_reg = None
        for W in mdl.parameters():
                if W.ndimension() < 2:
                        continue
                else:   
                        wt = torch.transpose(W,0,1)
                        m  = torch.matmul(wt,W)
                        ident = Variable(torch.eye(cols,cols),requires_grad=True)
                        ident = ident.cuda()
                        w_tmp = (m - ident)
                        if l2_reg is None:
                                l2_reg = ((torch.max(torch.abs(w_tmp)))**2)
                        else:   
                                l2_reg = l2_reg + ((torch.max(torch.abs(w_tmp)))**2)
                        #call("nvidia-smi")
        return l2_reg

def train(train_loader, model, criterion, optimizer, epoch,odecay):
    """Train for one epoch on the training set"""
    for i, (input, target) in enumerate(train_loader):
        target = target.cuda(async=True)
        input = input.cuda()
        input_var = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)

        # compute output
        output = model(input_var)

        #Compute Loss
        oloss =  l2_rego(model)
        loss = criterion(output, target_var)
        loss = loss + oloss

loss and oloss here is of type Variable. Within the first epoch I am getting
memory error.RuntimeError: cuda runtime error (2) : out of memory at /pytorch/torch/lib/THC/generic/THCStorage.cu:58
I fear, I am also hitting the same issue, but not able to quite deduce.