Spaces:
Runtime error
Runtime error
Gagan Bhatia
commited on
Commit
•
7ce22c0
1
Parent(s):
bd81828
Update train_model.py
Browse files
src/models/train_model.py
CHANGED
@@ -29,6 +29,9 @@ def train_model():
|
|
29 |
train_df=train_df,
|
30 |
eval_df=eval_df,
|
31 |
batch_size=params["batch_size"],
|
|
|
|
|
|
|
32 |
|
33 |
model.save_model(model_dir=params['model_dir'])
|
34 |
|
|
|
29 |
train_df=train_df,
|
30 |
eval_df=eval_df,
|
31 |
batch_size=params["batch_size"],
|
32 |
+
max_epochs=params["epochs"],
|
33 |
+
use_gpu=params["use_gpu"],
|
34 |
+
learning_rate=float(params["learning_rate"]),
|
35 |
|
36 |
model.save_model(model_dir=params['model_dir'])
|
37 |
|