forked from StrafesNET/strafe-ai
write model name based on num params
This commit is contained in:
@@ -160,7 +160,8 @@ fn training(gpu_id: usize, epochs: usize, learning_rate: f64) {
|
|||||||
let device = burn::backend::cuda::CudaDevice::new(gpu_id);
|
let device = burn::backend::cuda::CudaDevice::new(gpu_id);
|
||||||
|
|
||||||
let mut model: Net<TrainingBackend> = Net::init(&device);
|
let mut model: Net<TrainingBackend> = Net::init(&device);
|
||||||
println!("Training model ({} parameters)", model.num_params());
|
let num_params = model.num_params();
|
||||||
|
println!("Training model ({} parameters)", num_params);
|
||||||
|
|
||||||
let mut optim = AdamConfig::new().init();
|
let mut optim = AdamConfig::new().init();
|
||||||
|
|
||||||
@@ -205,7 +206,7 @@ fn training(gpu_id: usize, epochs: usize, learning_rate: f64) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
let date_string = format!("{}_{}.model", chrono::Utc::now(), best_loss);
|
let date_string = format!("{}_{}.model", num_params, best_loss);
|
||||||
best_model
|
best_model
|
||||||
.save_file(
|
.save_file(
|
||||||
date_string,
|
date_string,
|
||||||
|
|||||||
Reference in New Issue
Block a user