21 lines
518 B
Python
21 lines
518 B
Python
""" Module for PINA Optimizer """
|
|
|
|
import torch
|
|
from .utils import check_consistency
|
|
|
|
class Optimizer: # TODO improve interface
|
|
pass
|
|
|
|
|
|
class TorchOptimizer(Optimizer):
|
|
|
|
def __init__(self, optimizer_class, **kwargs):
|
|
check_consistency(optimizer_class, torch.optim.Optimizer, subclass=True)
|
|
|
|
self.optimizer_class = optimizer_class
|
|
self.kwargs = kwargs
|
|
|
|
def hook(self, parameters):
|
|
self.optimizer_instance = self.optimizer_class(
|
|
parameters, **self.kwargs
|
|
) |