diff --git a/image-inpainting/results/best_model.pt b/image-inpainting/results/best_model.pt
index a85a5e3..84b2251 100644
Binary files a/image-inpainting/results/best_model.pt and b/image-inpainting/results/best_model.pt differ
diff --git a/image-inpainting/src/__pycache__/architecture.cpython-313.pyc b/image-inpainting/src/__pycache__/architecture.cpython-313.pyc
new file mode 100644
index 0000000..dd2aca3
Binary files /dev/null and b/image-inpainting/src/__pycache__/architecture.cpython-313.pyc differ
diff --git a/image-inpainting/src/__pycache__/datasets.cpython-313.pyc b/image-inpainting/src/__pycache__/datasets.cpython-313.pyc
new file mode 100644
index 0000000..c101f32
Binary files /dev/null and b/image-inpainting/src/__pycache__/datasets.cpython-313.pyc differ
diff --git a/image-inpainting/src/__pycache__/train.cpython-313.pyc b/image-inpainting/src/__pycache__/train.cpython-313.pyc
new file mode 100644
index 0000000..f5138a9
Binary files /dev/null and b/image-inpainting/src/__pycache__/train.cpython-313.pyc differ
diff --git a/image-inpainting/src/__pycache__/utils.cpython-313.pyc b/image-inpainting/src/__pycache__/utils.cpython-313.pyc
new file mode 100644
index 0000000..215c119
Binary files /dev/null and b/image-inpainting/src/__pycache__/utils.cpython-313.pyc differ
diff --git a/image-inpainting/src/main.py b/image-inpainting/src/main.py
index 967a616..34f8f67 100644
--- a/image-inpainting/src/main.py
+++ b/image-inpainting/src/main.py
@@ -32,7 +32,7 @@ if __name__ == '__main__':
 
     config_dict['print_train_stats_at'] = 10
     config_dict['print_stats_at'] = 100
-    config_dict['plot_at'] = 10
+    config_dict['plot_at'] = 100
     config_dict['validate_at'] = 100
 
     network_config = {
diff --git a/image-inpainting/src/utils.py b/image-inpainting/src/utils.py
index e8e91f5..0a6f807 100644
--- a/image-inpainting/src/utils.py
+++ b/image-inpainting/src/utils.py
@@ -113,7 +113,7 @@ def create_predictions(model_config, state_dict_path, testset_path, device, save
             print(f"Processing image {i + 1}/{len(input_arrays)}")
             input_array = torch.from_numpy(input_arrays[i]).to(
                 device)
-            output = model(input_array)
+            output = model(input_array.unsqueeze(0) if hasattr(input_array, 'dim') and input_array.dim() == 3 else input_array)
             output = output.cpu().numpy()
             predictions.append(output)