Skip to content
Snippets Groups Projects
Commit 540b0bc3 authored by Alexandre Chapin's avatar Alexandre Chapin :race_car:
Browse files

Change Trainer

parent d8668224
No related branches found
No related tags found
No related merge requests found
......@@ -178,7 +178,7 @@ if __name__ == '__main__':
# Intialize training
params = [p for p in model.parameters() if p.requires_grad] # only keep trainable parameters
optimizer = optim.Adam(params, lr=lr_scheduler.get_cur_lr(0)) # to check after bnb.optim.Adam8bit(params, lr=lr_scheduler.get_cur_lr(0))
trainer = OSRTSamTrainer(model, optimizer, cfg, device, out_dir, train_dataset.render_kwargs)
trainer = SRTTrainer(model, optimizer, cfg, device, out_dir, train_dataset.render_kwargs)
checkpoint = Checkpoint(out_dir, device=device, encoder=encoder_module,
decoder=decoder_module, optimizer=optimizer)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment