diff --git a/test/TensorFlowNET.Keras.UnitTest/Gradient.cs b/test/TensorFlowNET.Keras.UnitTest/GradientTest.cs
similarity index 100%
rename from test/TensorFlowNET.Keras.UnitTest/Gradient.cs
rename to test/TensorFlowNET.Keras.UnitTest/GradientTest.cs
diff --git a/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelLoad.cs b/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelLoad.cs
index 385ec0f7..73f99bbc 100644
--- a/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelLoad.cs
+++ b/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelLoad.cs
@@ -44,7 +44,7 @@ public class SequentialModelLoad
         {
             TrainDir = "mnist",
             OneHot = false,
-            ValidationSize = 50000,
+            ValidationSize = 58000,
         }).Result;
 
         model.fit(dataset.Train.Data, dataset.Train.Labels, batch_size, num_epochs);
diff --git a/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelSave.cs b/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelSave.cs
index 1cf68d3b..e68fa9b4 100644
--- a/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelSave.cs
+++ b/test/TensorFlowNET.Keras.UnitTest/SaveModel/SequentialModelSave.cs
@@ -18,15 +18,15 @@ public class SequentialModelSave
     [TestMethod]
     public void SimpleModelFromAutoCompile()
     {
-        var inputs = tf.keras.layers.Input((28, 28, 1));
-        var x = tf.keras.layers.Flatten().Apply(inputs);
-        x = tf.keras.layers.Dense(100, activation: tf.nn.relu).Apply(x);
-        x = tf.keras.layers.Dense(units: 10).Apply(x);
-        var outputs = tf.keras.layers.Softmax(axis: 1).Apply(x);
-        var model = tf.keras.Model(inputs, outputs);
+        var inputs = keras.layers.Input((28, 28, 1));
+        var x = keras.layers.Flatten().Apply(inputs);
+        x = keras.layers.Dense(100, activation: tf.nn.relu).Apply(x);
+        x = keras.layers.Dense(units: 10).Apply(x);
+        var outputs = keras.layers.Softmax(axis: 1).Apply(x);
+        var model = keras.Model(inputs, outputs);
 
         model.compile(new Adam(0.001f), 
-            tf.keras.losses.SparseCategoricalCrossentropy(), 
+            keras.losses.SparseCategoricalCrossentropy(), 
             new string[] { "accuracy" });
 
         var data_loader = new MnistModelLoader();
@@ -37,7 +37,7 @@ public class SequentialModelSave
         {
             TrainDir = "mnist",
             OneHot = false,
-            ValidationSize = 10000,
+            ValidationSize = 58000,
         }).Result;
 
         model.fit(dataset.Train.Data, dataset.Train.Labels, batch_size, num_epochs);
@@ -69,7 +69,7 @@ public class SequentialModelSave
         {
             TrainDir = "mnist",
             OneHot = false,
-            ValidationSize = 50000,
+            ValidationSize = 58000,
         }).Result;
 
         model.fit(dataset.Train.Data, dataset.Train.Labels, batch_size, num_epochs);