fix

Signed-off-by: Zhi Lin <[email protected]>
oap-project · Apr 10, 2024 · cc18432 · cc18432
1 parent eff6b64
commit cc18432
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/python/raydp/tests/test_tf.py b/python/raydp/tests/test_tf.py
@@ -45,7 +45,7 @@ def test_tf_estimator(spark_on_ray_small, use_fs_directory):
     # create model
     model = keras.Sequential(
         [
-            keras.layers.InputLayer(input_shape=(2,)),
+            keras.layers.InputLayer(input_shape=()),
             keras.layers.Dense(1, activation='sigmoid')
         ]
     )

diff --git a/python/raydp/torch/estimator.py b/python/raydp/torch/estimator.py
@@ -392,4 +392,4 @@ def fit_on_spark(self,
 
     def get_model(self):
         assert self._trainer is not None, "Must call fit first"
-        return TorchCheckpoint(self._trained_results.checkpoint.to_directory()).get_model()
+        return TorchCheckpoint(self._trained_results.checkpoint.to_directory()).get_model(self._model)