Release GPU memory after tensorflow session is closed

run_release_gpu.py

def run_release_gpu(func):
  def parallel_wrapper(output_dict, *argv, **kwargs):
    ret = func(*argv, **kwargs)
    if ret is not None:
      output_dict['ret'] = ret
  
  def outer_wrapper(*argv, **kwargs):
    same_process = kwargs.pop('same_process', False)
    if same_process:
      return func(*argv, **kwargs)

    with multiprocessing.Manager() as manager:
      output = manager.dict()
      args = (output, ) + argv
      p = multiprocessing.Process(target=parallel_wrapper, args=args, kwargs=kwargs)
      p.start()
      p.join()
      ret_val = output.get('ret', None)

    return ret_val
  
  return outer_wrapper


@run_release_gpu
def run_computations(a, b, c):
  with tf.Graph().as_default():
    with tf.Session() as sess:
      pass
    
    
a = run_computations(1, 2, 3)

It's being a while since I posted this, no idea whether it is still a viable workaround.