Ok, I see. I just skimmed your repo and tried to locate the plf functions.
I created a small example, which leads to a memory issue (also on CPU).

class PartialLoss(object):
    """ Partially applied loss object. Has forward and zero_grad methods """
    def __init__(self):
        self.nets = []

class PartialXentropy(PartialLoss):
    def __init__(self, classifier):
        super(PartialXentropy, self).__init__()
        self.classifier = classifier
        self.nets.append(self.classifier)    


class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc1 = nn.Linear(1000, 10000)
        self.fc2 = nn.Linear(10000, 10000)
        
    def forward(self, x):
        x = self.fc1(x)
        x = F.relu(x)
        x = self.fc2(x)
        x = F.relu(x)
        return x

x = Variable(torch.randn(1, 1000))
model = MyModel()

output = model(x)

part_loss = PartialXentropy(model)

del model
del part_loss

Could you run this code and check your memory?
Even if I delete the model and the part_loss, which holds a reference to the model, I won’t get the memory back. Can you reproduce this?
EDIT: It seems the memory is just peaking and released afterwards. Sorry for the misleading idea.