Method for recording batches if not already recorded?

[simonlevine]

Hi, I'm trying to use this framework at the batch level within another training framework. Is there a way to make it possible to try using a recorded output before recording? Or to record only if the saved record hasn't yet been recorded?

something like:

from s2re import CachingHook

hook = CachingHook(path='/tmp/r3', backend='leveldb')

try:
  with hook.use(model):
      output = model(**batch)
      ...

 except:
    with hook.record(model):
        output = model(**batch)
        ...

The above doesn't work as hook.use establishes the lock in the directory. Thanks!