diff --git a/constants.py b/constants.py
index ca2dd8a..cbb837d 100644
--- a/constants.py
+++ b/constants.py
@@ -35,9 +35,9 @@
     "PCA_INPUT_SIZE": 8 + (2 * 10),  # PCA_ORTHOGONAL_BASES + (2 * H)
     "GAUSSIAN_MATRIX_FILE_PATH": "data/gaussian_matrix.npz",
     # ---  NEURAL MODEL CONSTANTS
-    "NEURAL_BATCH_SIZE": 2048 * 1,
+    "NEURAL_BATCH_SIZE": 256 * 4,
     "NEURAL_LEARNING_RATE": 0.01,
-    "NEURAL_N_EPOCHS": 30,
+    "NEURAL_N_EPOCHS": 50,
     "NEURAL_H": 12,
     "NEURAL_SIGMA_XY": 3.0,
     "NEURAL_SIGMA_UV": 0.6,
diff --git a/neural_model.py b/neural_model.py
index 4a31b04..8cd528c 100644
--- a/neural_model.py
+++ b/neural_model.py
@@ -116,13 +116,9 @@ def __init__(self, gaussian_matrix_xy, gaussian_matrix_uv):
             nn.ReLU(),
             nn.Linear(128, 64),
             nn.ReLU(),
-            nn.Linear(64, 32),
+            nn.Linear(64, 64),
             nn.ReLU(),
-            nn.Linear(32, 16),
-            nn.ReLU(),
-            nn.Linear(16, 16),
-            nn.ReLU(),
-            nn.Linear(16, 1),
+            nn.Linear(64, 1),
         )
 
     def forward(self, x):