def _with_model(self, model: Model) -> TrainingLoop:
return self.cls(
model=model,
triples_factory=self.triples_factory,
automatic_memory_optimization=False,
optimizer=self.optimizer_cls(model.get_grad_params()),
)
def __init__(self, model: Model, patience: int, automatic_memory_optimization: bool = False):
super().__init__(
model=model,
optimizer=optim.Adam(lr=1.0, params=model.parameters()),
automatic_memory_optimization=automatic_memory_optimization,
)
self.patience = patience
def __init__(
self,
model: Model,
triples_factory: TriplesFactory,
sub_batch_size: int,
automatic_memory_optimization: bool = False,
):
super().__init__(
model=model,
triples_factory=triples_factory,
optimizer=optim.Adam(lr=1.0, params=model.parameters()),
automatic_memory_optimization=automatic_memory_optimization,
)
self.sub_batch_size = sub_batch_size