I am trying to run a small neural network on the CPU and am finding that the memory used by my script increases without limit. Since my script does not do much besides call the network, the problem appears to be a memory leak within pytorch.
The problem does not occur if I run the model on the gpu.
I’ve also posted this to the pytorch github, but I was hoping someone on here might be able to help.
Here’s the script I’m using. The output of the program increases until the machine runs out of memory, at which point it crashes.
import torch
import torch.nn.functional as F
import torch.nn as nn
import resource
import gc
class Net(nn.Module):
def init(self, in_size, h_size, out_size):
super(Net, self).init()
    self.lin1 = nn.Linear(in_size, h_size)
    self.lin2 = nn.Linear(h_size, out_size)
    self.nonlin = F.relu

def forward(self, x):
    h = self.nonlin(self.lin1(x))
    h = self.nonlin(self.lin2(h))
    return h

def main():
in_size = 100
h_size = 256
out_size = 512
net = Net(in_size, h_size, out_size)

while True:
    gc.collect()
    print('maxrss = {}'.format(
        resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1e6))
    for i in range(1000):
        net(torch.randn(in_size))

if name == ‘main’:
main()