post_training_files

post_training_files.py

# below only necessary if we're using tracking
accelerator.end_training()
if accelerator.is_main_process:
    run_config_local_path = write_run_config(hyperparams['train_batch_size'], hyperparams['eval_batch_size'],
                                             hyperparams['learning_rate'], hyperparams['epochs'])
    local_best_model = os.path.join(os.getcwd(), f"{hyperparams['model_name']}_finetuned.pth")
    accelerator.save(best_model.state_dict(), local_best_model)
    local_torchscript_model_path = os.path.join(os.getcwd(), f"{hyperparams['model_name']}_finetuned.pt")
    convert_to_torchscript(local_torchscript_model_path,best_model, train_data, hyperparams)