Update app.py

app.py

@@ -131,6 +131,10 @@ def generate_video(
     else:
         baseline_video_path = None
 
+    if torch.cuda.is_available():
+        print("Allocated:", torch.cuda.memory_allocated() / 1024**2, "MB")
+        print("Cached:   ", torch.cuda.memory_reserved() / 1024**2, "MB")
+
     return nag_video_path, baseline_video_path, current_seed
 
 
@@ -147,6 +151,9 @@ def generate_video_with_example(
         seed=DEFAULT_SEED, randomize_seed=False,
         compare=True,
     )
+    if torch.cuda.is_available():
+        print("Allocated:", torch.cuda.memory_allocated() / 1024**2, "MB")
+        print("Cached:   ", torch.cuda.memory_reserved() / 1024**2, "MB")
     return nag_video_path, baseline_video_path, \
         DEFAULT_H_SLIDER_VALUE, DEFAULT_W_SLIDER_VALUE, \
         DEFAULT_DURATION_SECONDS, DEFAULT_STEPS, seed, True