def build_tft(scaler):
    tft = TFTModel(
        input_chunk_length=48,
        output_chunk_length=24,
        hidden_size=8,
        lstm_layers=3,
        num_attention_heads=4,
        full_attention=False,
        feed_forward='GatedResidualNetwork',
        dropout=0.15,
        loss_fn=QuantileRegression(quantiles=[0.1, 0.5, 0.9]),
        likelihood=QuantileRegression(quantiles=[0.1, 0.5, 0.9]),
        use_static_covariates=True,
        torch_metrics=MetricCollection(
            metrics=[MeanAbsoluteError(), MeanAbsolutePercentageError()]
        ),
        batch_size=1024,
        n_epochs=15,
        model_name=f"tft_{8}_{3}_{4}_{0.00003}",
        work_dir=os.path.join(args.save_dir, 'tft'),
        log_tensorboard=True,
        save_checkpoints=True,
        optimizer_kwargs={"lr": 0.00003},
        random_state=42,
        add_encoders={
            'cyclic': {'future': ['month']},
            'datetime_attribute': {'future': ['hour']},
            'transformer': scaler
        }
    )

    return tft

# scaler is a loaded darts.dataprocessing.transformers.Scaler saved during preprocessing
model = build_tft(scaler)
model.fit(trian, val)

model = TFTModel.load_from_checkpoint(model_name='tft_8_3_4_0.0003', work_dir='training_data/tft', best=True)