How can I free GPU memory allocated with NetTrain

Question

When I train a neural network with NetTrain[..., TargetDevice->"GPU"], wolfram's neural network code automatically allocates memory on my GPU for my model's computation graph, copied depending on the batch size.

My problem is that if I change my model and re-run NetTrain, memory allocated for the old model is not reliably de-allocated from my gpu, resulting in an error:

NetTrain: An unknown internal error occurred. Consult internal`$LastInternalFailure for potential information.

and internal`$LastInternalFailure of:

MXNetError <...> ./pooled_storage_manager.h:161: cudaMalloc retry failed: out of memory

How can I explicitly tell the wolfram neural network library to clear out allocated GPU memory?

(I have found a workaround without restarting my computer involving repeating restarting the Wolfram Kernel and waiting dozens of seconds until my gpu memory is freed. But this is extremely inconvenient, and it requires me to re-run my notebook to continue from where I left off.)

Answer 1

This issue appears to be resolved in Wolfram 13.1! (the memory is not freed by mxnet, but at least I can ClearAll the model and re-train and without the GPU memory consumption spiking).

Repeatedly running the following cell in a notebook does not result in GPU memory allocation errors anymore (note, this uses about 4 gb of VRAM for me, but you could reduce the number of layers if you want to test on a GPU with less VRAM).

$HistoryLength = 0;
ClearAll[data, trained, net]
data = Flatten@
   Table[{x, y} -> Exp[-Norm[{x, y}]], {x, -3, 3, .005}, {y, -3, 
     3, .005}];
net = NetChain[{300, Tanh, 3000, Sin, 3000, Sin, 3000, Sin, 3000, 
    Sin, 3000, Tanh, 1}];
trained = 
 NetTrain[net, data, BatchSize -> 1024, TargetDevice -> "GPU"]