From 5be474d178206a0e7da241cee261753827ffe390 Mon Sep 17 00:00:00 2001
From: z <5618955+lilzeta@users.noreply.github.com>
Date: Mon, 11 Nov 2024 14:44:13 -0800
Subject: [PATCH] Empty torch cache on sample button start

---
 scripts/demo/video_sampling.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/scripts/demo/video_sampling.py b/scripts/demo/video_sampling.py
index 1f4fcfc4a..73f19652c 100644
--- a/scripts/demo/video_sampling.py
+++ b/scripts/demo/video_sampling.py
@@ -2,6 +2,7 @@
 import sys
 
 sys.path.append(os.path.realpath(os.path.join(os.path.dirname(__file__), "../../")))
+import torch
 from pytorch_lightning import seed_everything
 from scripts.demo.streamlit_helpers import *
 from scripts.demo.sv3d_helpers import *
@@ -251,6 +252,7 @@
             saving_fps = value_dict["fps"]
 
         if st.button("Sample"):
+            torch.cuda.empty_cache()
             out = do_sample(
                 model,
                 sampler,