Fixed small summary with model architectures until now

2023-11-30 21:53:35 +00:00
parent eba10c8f83
commit 120b6aa5bd
23 changed files with 402 additions and 185 deletions
--- a/src/trainers/autoregressive_trainer.py
+++ b/src/trainers/autoregressive_trainer.py
@@ -15,6 +15,7 @@ class AutoRegressiveTrainer(Trainer):
    def __init__(
        self,
        model: torch.nn.Module,
+        input_dim: tuple,
        optimizer: torch.optim.Optimizer,
        criterion: torch.nn.Module,
        data_processor: DataProcessor,
@@ -23,6 +24,7 @@ class AutoRegressiveTrainer(Trainer):
    ):
        super().__init__(
            model=model,
+            input_dim=input_dim,
            optimizer=optimizer,
            criterion=criterion,
            data_processor=data_processor,