clear cache on all exceptions

This commit is contained in:
Lincoln Stein 2024-04-11 14:51:32 -04:00 committed by psychedelicious
parent 46d23cd868
commit 651c0b39b1

View File

@ -273,7 +273,7 @@ class ModelCache(ModelCacheBase[AnyModel]):
snapshot_before = self._capture_memory_snapshot()
try:
cache_entry.model.to(target_device)
except RuntimeError as e: # blow away cache entry
except Exception as e: # blow away cache entry
self._delete_cache_entry(cache_entry)
raise e