My guess is in the training loop you assign various variables, likely something like prediction = net(train_data). This prediction has pointers to the rest of the entire computation graph the data goes through (so probably a lot of memory). Once the training loop finishes, the last assignment to prediction is still kept, since in Python loops and ifs don’t create new scopes. All those variables are in the function scope (or global scope). So there’s no way to know that that whole graph should now be reclaimed.
Either explicitly del the variables after the training loop, or put the training loop into a separate function.
By they way, also make sure you’re not doing something like mean_err += err where err is a Variable, because that will also create a large unnecessary graph. Use err.data instead.