Update train.py

bumble-tech · Nov 5, 2023 · 3a5c81e · 3a5c81e
1 parent 33a87bf
commit 3a5c81e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -214,7 +214,7 @@ def train(train_id: str,
     logger.info(f'Training complete, saving model to {train_id}')
 
     model.save(train_id)
-    model.save_base64_serving(train_id, input_dtype=dtype)
+    model.save_base64_serving(f'{train_id}_base64', input_dtype=dtype)
 
     logger.info(f'Model saved to {train_id}')