mirror of
https://github.com/ostris/ai-toolkit.git
synced 2026-01-26 16:39:47 +00:00
Added Model rescale and prepared a release upgrade
This commit is contained in:
@@ -54,6 +54,8 @@ def get_optimizer(
|
||||
elif lower_type == 'lion':
|
||||
from lion_pytorch import Lion
|
||||
return Lion(params, lr=learning_rate, **optimizer_params)
|
||||
elif lower_type == 'adagrad':
|
||||
optimizer = torch.optim.Adagrad(params, lr=float(learning_rate), **optimizer_params)
|
||||
else:
|
||||
raise ValueError(f'Unknown optimizer type {optimizer_type}')
|
||||
return optimizer
|
||||
|
||||
Reference in New Issue
Block a user