From 8917255002914103956a6cab426a091ad9fc3b8a Mon Sep 17 00:00:00 2001
From: Sumeth Sathnindu <65599800+LucipherDev@users.noreply.github.com>
Date: Sun, 19 Jan 2025 10:40:49 +0530
Subject: [PATCH] Bug fix

Fixed error in catching OOM
---
 nodes.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/nodes.py b/nodes.py
index c2ea1de..d469ecb 100644
--- a/nodes.py
+++ b/nodes.py
@@ -273,7 +273,9 @@ class TangoFluxVAEDecodeAndPlay:
                 results = torch.cat(results, dim=0)
             return results
 
-        except torch.OutOfMemoryError:
+        except (torch.cuda.OutOfMemoryError, RuntimeError) as e:
+            if isinstance(e, RuntimeError) and "out of memory" not in str(e):
+                raise e
             torch.cuda.empty_cache()
             log.warning("OOM encountered. Falling back to tiled decoding.")
             return self.decode_tiled(vae, latents, tile_size)