Reduce memory requirements of float64 test.

james-choncholas · Jan 21, 2025 · c1bac5d · c1bac5d
1 parent afe5601
commit c1bac5d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tf_shell_ml/test/postscale_model_float64_test.py b/tf_shell_ml/test/postscale_model_float64_test.py
@@ -33,7 +33,7 @@ def _test_model(self, disable_encryption, disable_masking, disable_noise, cache)
 
         # Clip dataset images to limit memory usage. The model accuracy will be
         # bad but this test only measures functionality.
-        x_train, x_test = x_train[:, :350], x_test[:, :350]
+        x_train, x_test = x_train[:, :250], x_test[:, :250]
 
         labels_dataset = tf.data.Dataset.from_tensor_slices(y_train)
         labels_dataset = labels_dataset.batch(2**10)
@@ -84,7 +84,7 @@ def _test_model(self, disable_encryption, disable_masking, disable_noise, cache)
             validation_data=val_dataset,
         )
 
-        self.assertGreater(history.history["val_categorical_accuracy"][-1], 0.25)
+        self.assertGreater(history.history["val_categorical_accuracy"][-1], 0.20)
 
     def test_model(self):
         tf.keras.backend.set_floatx("float64")