code: def gpu_mem_restore(func):
      "Reclaim GPU RAM if CUDA out of memory happened, or execution was interrupted"
      @functools.wraps(func)
      def wrapper(*args, **kwargs):
          tb_clear_frames = os.environ.get('FASTAI_TB_CLEAR_FRAMES', None)
          if not IS_IN_IPYTHON or tb_clear_frames=="0":
              return func(*args, **kwargs)
  
          try:
              return func(*args, **kwargs)
          except Exception as e:
              if ("CUDA out of memory" in str(e) or
                  "device-side assert triggered" in str(e) or
                  tb_clear_frames == "1"):
                  type, val, tb = get_ref_free_exc_info() # must!
                  gc.collect()
                  if "device-side assert triggered" in str(e):
                      warn()
                  raise type(val).with_traceback(tb) from None
              else: raise # re-raises the exact last exception
      return wrapper
docs: Reclaim GPU RAM if CUDA out of memory happened, or execution was interrupted