change layout and fix issue with natural image

app.py

@@ -34,7 +34,7 @@ def generate_imgs_from_user(image,
                             physics: PhysicsWithGenerator, use_gen: bool,
                             metrics: List[Metric]):
     # PIL image -> torch.Tensor
-    x = transforms.ToTensor()(image).unsqueeze(0)
+    x = transforms.ToTensor()(image).unsqueeze(0).to('cuda')
 
     return generate_imgs(x, model, baseline, physics, use_gen, metrics)
     
@@ -147,7 +147,7 @@ with gr.Blocks(title=title, css=custom_css) as interface:
     # Loading things
     model_a_placeholder = gr.State(lambda: get_eval_model_on_DEVICE_STR("unext_emb_physics_config_C", ""))  # lambda expression to instanciate a callable in a gr.State
     model_b_placeholder = gr.State(lambda: get_baseline_model_on_DEVICE_STR("DPIR"))  # lambda expression to instanciate a callable in a gr.State
-    dataset_placeholder = gr.State(get_dataset_on_DEVICE_STR("Natural"))
+    dataset_placeholder = gr.State(lambda: get_dataset("Natural"))
     physics_placeholder = gr.State(lambda: get_physics_on_DEVICE_STR("MotionBlur_easy"))  # lambda expression to instanciate a callable in a gr.State
     metrics_placeholder = gr.State(get_list_metrics_on_DEVICE_STR(["PSNR"]))
 
@@ -196,12 +196,12 @@ with gr.Blocks(title=title, css=custom_css) as interface:
 
         # Components: Load Metric + Load image Buttons
         with gr.Row():
-            with gr.Column(scale=2):
+            with gr.Column(scale=3):
                 choose_metrics = gr.CheckboxGroup(choices=Metric.all_metrics,
                                                   value=metric_names,
                                                   label="Choose metrics you are interested")
             use_generator_button = gr.Checkbox(label="Generate valid physics parameters", scale=1)
-            run_button = gr.Button("Run current image")
+            run_button = gr.Button("Run current image", scale=1)
             with gr.Column(scale=1):
                 load_button = gr.Button("Load images from dataset...")
                 load_random_button = gr.Button("Load randomly from dataset...")