From 5b36bcd1eec5e101f4b305a478ba33ef200c5eb1 Mon Sep 17 00:00:00 2001 From: Arne Binder Date: Fri, 24 Jan 2025 23:53:18 +0100 Subject: [PATCH] override to() to call it on base model --- src/pie_modules/models/simple_generative.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/pie_modules/models/simple_generative.py b/src/pie_modules/models/simple_generative.py index a4bdae0b8..75b43fab9 100644 --- a/src/pie_modules/models/simple_generative.py +++ b/src/pie_modules/models/simple_generative.py @@ -9,7 +9,7 @@ from torch.optim import Optimizer from transformers import PreTrainedModel, get_linear_schedule_with_warmup from transformers.modeling_outputs import Seq2SeqLMOutput -from typing_extensions import TypeAlias +from typing_extensions import Self, TypeAlias from pie_modules.models.common import ModelWithBoilerplate from pie_modules.utils import resolve_type @@ -159,3 +159,7 @@ def configure_optimizers(self) -> OptimizerLRScheduler: return [optimizer], [{"scheduler": scheduler, "interval": "step"}] else: return optimizer + + def to(self, *args: Any, **kwargs: Any) -> Self: + self.model = self.model.to(*args, **kwargs) + return self