debug print

app.py

@@ -62,6 +62,7 @@ def generate_imgs(x: torch.Tensor,
                   physics: PhysicsWithGenerator, use_gen: bool,
                   baseline: BaselineModel, model: EvalModel,
                   metrics: List[Metric]):
+    print(torch.cuda.memory_allocated() / 1024**2)
     ### Compute y
     y = physics(x, use_gen)  # possible reduction in img shape due to Blurring
 
@@ -113,6 +114,7 @@ def generate_imgs(x: torch.Tensor,
     out = to_pil(process_img(out)[0].to('cpu'))
     out_baseline = to_pil(process_img(out_baseline)[0].to('cpu'))
 
+    print(torch.cuda.memory_allocated() / 1024**2)
     return x, y, out, out_baseline, physics.display_saved_params(), metrics_y, metrics_out, metrics_out_baseline
 
 
@@ -150,7 +152,7 @@ generate_imgs_from_user_partial = partial(generate_imgs_from_user, model=ram_mod
 generate_imgs_from_dataset_partial = partial(generate_imgs_from_dataset, model=ram_model, metrics=psnr)
 generate_random_imgs_from_dataset_partial = partial(generate_random_imgs_from_dataset, model=ram_model, metrics=psnr)
 
-
+print(torch.cuda.memory_allocated() / 1024**2)
 ### Gradio Blocks interface
 
 title = "Inverse problem playground"  # displayed on gradio tab and in the gradio page
@@ -166,7 +168,7 @@ with gr.Blocks(title=title, theme=gr.themes.Glass()) as interface:
     physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))
     model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))
 
-
+    print(torch.cuda.memory_allocated() / 1024**2)
     @gr.render(inputs=[dataset_placeholder, physics_placeholder, available_physics_placeholder])
     def dynamic_layout(dataset, physics, available_physics):
         ### LAYOUT