This commit is contained in:
Nicola Demo
2024-08-05 17:34:34 +02:00
parent 686b557144
commit 5245a0b68c
19 changed files with 483 additions and 173 deletions

View File

@@ -0,0 +1,19 @@
""" Module for PINA Torch Optimizer """
import torch
from ..utils import check_consistency
from .optimizer_interface import Optimizer
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
)