File size: 252 Bytes
78a947a
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
"""
Contains a function to return the Adam optimizer.

Author: Ole-Christian Galbo Engstrøm
E-mail: ocge@foss.dk
"""

import torch


def Adam(model: torch.nn.Module, lr: float):
    optim = torch.optim.Adam(model.parameters(), lr=lr)
    return optim