diff --git a/notebooks/finetuning_example.py b/notebooks/finetuning_example.py
new file mode 100644
index 0000000..5664456
--- /dev/null
+++ b/notebooks/finetuning_example.py
@@ -0,0 +1,224 @@
+"""
+Example usage of the TimesFM Finetuning Framework.
+"""
+
+import yfinance as yf
+import torch
+from os import path
+import numpy as np
+from torch.utils.data import Dataset
+from timesfm import TimesFm, TimesFmHparams, TimesFmCheckpoint
+from timesfm.pytorch_patched_decoder import PatchedTimeSeriesDecoder
+from finetuning_torch import FinetuningConfig, TimesFMFinetuner
+from huggingface_hub import snapshot_download
+import numpy as np
+import pandas as pd
+from torch.utils.data import Dataset
+import torch
+import yfinance as yf
+from typing import Tuple, Optional
+
+from timesfm import TimesFm, TimesFmHparams
+
+
+class TimeSeriesDataset(Dataset):
+    """Dataset for time series data compatible with TimesFM."""
+
+    def __init__(self, series: np.ndarray, context_length: int, horizon_length: int):
+        """
+        Initialize dataset.
+
+        Args:
+            series: Time series data
+            context_length: Number of past timesteps to use as input
+            horizon_length: Number of future timesteps to predict
+        """
+        self.series = series
+        self.context_length = context_length
+        self.horizon_length = horizon_length
+        self._prepare_samples()
+
+    def _prepare_samples(self) -> None:
+        """Prepare sliding window samples from the time series."""
+        self.samples = []
+        total_length = self.context_length + self.horizon_length
+
+        for start_idx in range(0, len(self.series) - total_length + 1):
+            end_idx = start_idx + self.context_length
+            x_context = self.series[start_idx:end_idx]
+            x_future = self.series[end_idx : end_idx + self.horizon_length]
+            self.samples.append((x_context, x_future))
+
+    def __len__(self) -> int:
+        return len(self.samples)
+
+    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        x_context, x_future = self.samples[index]
+
+        x_context = torch.tensor(x_context, dtype=torch.float32)
+        x_future = torch.tensor(x_future, dtype=torch.float32)
+
+        input_padding = torch.zeros_like(x_context)
+        freq = torch.zeros(1, dtype=torch.long)
+
+        return x_context, input_padding, freq, x_future
+
+
+def prepare_datasets(
+    series: np.ndarray, context_length: int, horizon_length: int, train_split: float = 0.8
+) -> Tuple[Dataset, Dataset]:
+    """
+    Prepare training and validation datasets from time series data.
+
+    Args:
+        series: Input time series data
+        context_length: Number of past timesteps to use
+        horizon_length: Number of future timesteps to predict
+        train_split: Fraction of data to use for training
+
+    Returns:
+        Tuple of (train_dataset, val_dataset)
+    """
+    train_size = int(len(series) * train_split)
+    train_data = series[:train_size]
+    val_data = series[train_size:]
+
+    # Create datasets
+    train_dataset = TimeSeriesDataset(train_data, context_length=context_length, horizon_length=horizon_length)
+
+    val_dataset = TimeSeriesDataset(val_data, context_length=context_length, horizon_length=horizon_length)
+
+    return train_dataset, val_dataset
+
+
+def get_model(load_weights: bool = False):
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    repo_id = "google/timesfm-2.0-500m-pytorch"
+    hparams = TimesFmHparams(
+        backend=device,
+        per_core_batch_size=32,
+        horizon_len=128,
+        num_layers=50,
+        use_positional_embedding=False,
+        context_len=192,
+    )
+    tfm = TimesFm(hparams=hparams, checkpoint=TimesFmCheckpoint(huggingface_repo_id=repo_id))
+
+    model = PatchedTimeSeriesDecoder(tfm._model_config)
+    if load_weights:
+        checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
+        loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
+        model.load_state_dict(loaded_checkpoint)
+    model = model.to(device)
+    return model, hparams, tfm._model_config
+
+
+def plot_predictions(
+    model: TimesFm,
+    val_dataset: Dataset,
+    save_path: Optional[str] = "predictions.png",
+) -> None:
+    """
+    Plot model predictions against ground truth for a batch of validation data.
+
+    Args:
+        model: Trained TimesFM model
+        val_dataset: Validation dataset
+        save_path: Path to save the plot
+    """
+    import matplotlib.pyplot as plt
+
+    model.eval()
+
+    x_context, x_padding, freq, x_future = val_dataset[0]
+    x_context = x_context.unsqueeze(0)  # Add batch dimension
+    x_padding = x_padding.unsqueeze(0)
+    freq = freq.unsqueeze(0)
+    x_future = x_future.unsqueeze(0)
+
+    device = next(model.parameters()).device
+    x_context = x_context.to(device)
+    x_padding = x_padding.to(device)
+    freq = freq.to(device)
+    x_future = x_future.to(device)
+
+    with torch.no_grad():
+        predictions = model(x_context, x_padding.float(), freq)
+        predictions_mean = predictions[..., 0]  # [B, N, horizon_len]
+        last_patch_pred = predictions_mean[:, -1, :]  # [B, horizon_len]
+
+    context_vals = x_context[0].cpu().numpy()
+    future_vals = x_future[0].cpu().numpy()
+    pred_vals = last_patch_pred[0].cpu().numpy()
+
+    context_len = len(context_vals)
+    horizon_len = len(future_vals)
+
+    plt.figure(figsize=(12, 6))
+
+    plt.plot(range(context_len), context_vals, label="Historical Data", color="blue", linewidth=2)
+
+    plt.plot(
+        range(context_len, context_len + horizon_len),
+        future_vals,
+        label="Ground Truth",
+        color="green",
+        linestyle="--",
+        linewidth=2,
+    )
+
+    plt.plot(range(context_len, context_len + horizon_len), pred_vals, label="Prediction", color="red", linewidth=2)
+
+    plt.xlabel("Time Step")
+    plt.ylabel("Value")
+    plt.title("TimesFM Predictions vs Ground Truth")
+    plt.legend()
+    plt.grid(True)
+
+    if save_path:
+        plt.savefig(save_path)
+        print(f"Plot saved to {save_path}")
+
+    plt.close()
+
+
+def get_data(context_len: int, horizon_len: int) -> Tuple[Dataset, Dataset]:
+    df = yf.download("AAPL", start="2010-01-01", end="2019-01-01")
+    time_series = df["Close"].values
+
+    train_dataset, val_dataset = prepare_datasets(
+        series=time_series,
+        context_length=context_len,
+        horizon_length=horizon_len,
+        train_split=0.8,
+    )
+
+    print(f"Created datasets:")
+    print(f"- Training samples: {len(train_dataset)}")
+    print(f"- Validation samples: {len(val_dataset)}")
+    return train_dataset, val_dataset
+
+
+def basic_example():
+    """Basic example of finetuning TimesFM on stock data."""
+    model, hparams, tfm_config = get_model(load_weights=True)
+    config = FinetuningConfig(batch_size=256, num_epochs=5, learning_rate=1e-4, use_wandb=False)
+
+    train_dataset, val_dataset = get_data(128, tfm_config.horizon_len)
+    finetuner = TimesFMFinetuner(model, config)
+
+    print("\nStarting finetuning...")
+    results = finetuner.finetune(train_dataset=train_dataset, val_dataset=val_dataset)
+
+    print("\nFinetuning completed!")
+    print(f"Training history: {len(results['history']['train_loss'])} epochs")
+
+    plot_predictions(
+        model=model,
+        val_dataset=val_dataset,
+        save_path="timesfm_predictions.png",
+    )
+
+
+if __name__ == "__main__":
+    basic_example()
diff --git a/notebooks/finetuning_torch.py b/notebooks/finetuning_torch.py
index 50a7f4d..bf908e7 100644
--- a/notebooks/finetuning_torch.py
+++ b/notebooks/finetuning_torch.py
@@ -1,274 +1,199 @@
-# Filename: tutorial_timesfm.py
+"""
+TimesFM Finetuner: A flexible framework for finetuning TimesFM models on custom datasets.
+
+Example usage:
+    ```python
+    # Prepare datasets
+    train_dataset = TimeSeriesDataset(train_data, context_length=128, horizon_length=32)
+    val_dataset = TimeSeriesDataset(val_data, context_length=128, horizon_length=32)
+    
+    # Initialize model and configuration
+    model = TimesFm(...)
+    config = FinetuningConfig(
+        batch_size=64,
+        num_epochs=50,
+        learning_rate=1e-4,
+        use_wandb=True
+    )
+    
+    # Create finetuner
+    finetuner = TimesFMFinetuner(model, config)
+    
+    # Finetune model
+    results = finetuner.finetune(train_dataset, val_dataset)
+    ```
+"""
+
+import abc
+import logging
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Optional, Dict, Any
 
-import yfinance as yf
-import numpy as np
-import pandas as pd
 import torch
 from torch.utils.data import Dataset, DataLoader
 import torch.optim as optim
-import timesfm
-from os import path
-from typing import Any, Sequence
+from torch.nn.parallel import DistributedDataParallel as DDP
+import wandb
+import multiprocessing as mp
 
-import numpy as np
-import torch
-from huggingface_hub import snapshot_download
+from timesfm import TimesFm
 
 
-from timesfm.pytorch_patched_decoder import TimesFMConfig, PatchedTimeSeriesDecoder
+@dataclass
+class FinetuningConfig:
+    """Configuration for TimesFM finetuning process."""
 
-import torch
-import matplotlib.pyplot as plt
+    # Training parameters
+    batch_size: int = 32
+    num_epochs: int = 20
+    learning_rate: float = 1e-4
+    weight_decay: float = 0.01
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Hardware parameters
+    device: str = "cuda" if torch.cuda.is_available() else "cpu"
+    distributed: bool = False
+    world_size: int = 1
+
+    # Logging parameters
+    use_wandb: bool = False
+    wandb_project: str = "timesfm-finetuning"
 
 
-# --------------------------------------------------
-# 1. Download stock data via yfinance
-# --------------------------------------------------
-def download_yfinance_data(ticker="AAPL", start="2020-01-01", end="2022-01-01"):
-    """
-    Download daily stock data for a given ticker from Yahoo Finance.
-    Returns a pandas DataFrame with columns like 'Open', 'High', 'Low', 'Close', 'Volume'.
-    """
-    df = yf.download(ticker, start=start, end=end)
-    df = df.dropna()
-    return df["Close"].reset_index(drop=True)
+class TimesFMFinetuner:
+    """Main class for finetuning TimesFM models."""
 
-
-# --------------------------------------------------
-# 2. Create a dataset class for TimesFM
-# --------------------------------------------------
-class FinancialDataset(Dataset):
     def __init__(
         self,
-        series: pd.Series,
-        config: TimesFMConfig,
-        context_length=128,  # how many past timesteps as input
-        horizon_length=32,  # how many future steps to predict
+        model: TimesFm,
+        config: FinetuningConfig,
+        loss_fn: Optional[callable] = None,
+        logger: Optional[logging.Logger] = None,
     ):
-        super().__init__()
+        """
+        Initialize TimesFM finetuner.
 
-        self.series = series.values.astype(np.float32)
-        self.context_length = context_length
-        self.horizon_length = horizon_length
+        Args:
+            model: TimesFM model to finetune
+            config: Finetuning configuration
+            logger: Optional logger instance
+        """
+        self.model = model
         self.config = config
+        self.logger = logger or logging.getLogger(__name__)
 
-        self.samples = []
-        # We want to ensure we have at least context_length + horizon_length points.
-        for start_idx in range(0, len(self.series) - (context_length + horizon_length)):
-            end_idx = start_idx + context_length
-            # context slice
-            x_context = self.series[start_idx:end_idx]
-            # future/horizon slice
-            x_future = self.series[end_idx : end_idx + horizon_length]
-            self.samples.append((x_context, x_future))
+        self.device = torch.device(config.device)
+        self.loss_fn = loss_fn or (lambda x, y: torch.mean((x - y.squeeze(-1)) ** 2))  # MSELoss()
 
-    def __len__(self):
-        return len(self.samples)
+        if config.use_wandb:
+            self._setup_wandb()
 
-    def __getitem__(self, index):
-        x_context, x_future = self.samples[index]
-        # Convert to torch
-        x_context = torch.tensor(x_context, dtype=torch.float32)
-        x_future = torch.tensor(x_future, dtype=torch.float32)
+    def _setup_wandb(self) -> None:
+        """Initialize Weights & Biases logging."""
+        wandb.init(project=self.config.wandb_project, entity=self.config.wandb_entity, config=self.config.__dict__)
 
-        input_padding = torch.zeros_like(x_context)
+    def _create_dataloader(self, dataset: Dataset, name: str) -> DataLoader:
+        """Create a dataloader from a dataset."""
+        return DataLoader(
+            dataset,
+            batch_size=self.config.batch_size,
+            shuffle=name == "train",
+            num_workers=mp.cpu_count(),
+            pin_memory=self.device.type == "cuda",
+            persistent_workers=True,
+            prefetch_factor=2,
+        )
 
-        freq = torch.zeros(1, dtype=torch.long)
+    def _train_epoch(self, train_loader: DataLoader, optimizer: torch.optim.Optimizer) -> float:
+        """Train for one epoch."""
+        self.model.train()
+        total_loss = 0.0
+        n_batches = len(train_loader)
 
-        return x_context, input_padding, freq, x_future
+        for batch in train_loader:
+            x_context, x_padding, freq, x_future = [t.to(self.device, non_blocking=True) for t in batch]
 
-
-def collate_fn(batch):
-    xs_context = [item[0] for item in batch]
-    xs_padding = [item[1] for item in batch]
-    freqs = [item[2] for item in batch]
-    xs_future = [item[3] for item in batch]
-
-    x_context = torch.stack(xs_context, dim=0)
-    input_pad = torch.stack(xs_padding, dim=0)
-    freq = torch.stack(freqs, dim=0)  # shape [B, 1]
-    x_future = torch.stack(xs_future, dim=0)
-
-    return x_context, input_pad, freq, x_future
-
-
-def get_model(*, load_weights: bool = False):
-    # standard model hack
-    repo_id = "google/timesfm-2.0-500m-pytorch"
-    tfm = timesfm.TimesFm(
-        hparams=timesfm.TimesFmHparams(
-            backend="cuda",
-            per_core_batch_size=32,
-            horizon_len=128,
-            num_layers=50,
-            use_positional_embedding=False,
-            context_len=192,
-        ),
-        checkpoint=timesfm.TimesFmCheckpoint(huggingface_repo_id=repo_id),
-    )
-
-    model = PatchedTimeSeriesDecoder(tfm._model_config)
-
-    if load_weights:
-        checkpoint_path = path.join(snapshot_download(repo_id), "torch_model.ckpt")
-        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
-        loaded_checkpoint = torch.load(checkpoint_path, weights_only=True)
-        model.load_state_dict(loaded_checkpoint)
-        print("After loading:")
-        print(model.state_dict()["input_ff_layer.hidden_layer.0.weight"])
-    model = model.to(device)
-
-    # import sys
-    # sys.exit(-1)
-    # repo_id = "google/timesfm-1.0-200m"
-    return model, tfm._model_config
-
-
-def train_model(
-    ticker="AAPL", start="2015-01-01", end="2022-01-01", train_split=0.8, batch_size=8, num_epochs=20, pretrained=False
-):
-    df_close = download_yfinance_data(ticker, start=start, end=end)
-    model, config = get_model(load_weights=pretrained)
-
-    total_len = len(df_close)
-    train_size = int(total_len * train_split)
-    val_size = total_len - train_size
-
-    train_series = df_close.iloc[:train_size].reset_index(drop=True)
-    val_series = df_close.iloc[train_size:].reset_index(drop=True)
-
-    train_dataset = FinancialDataset(
-        series=train_series, config=config, context_length=128, horizon_length=config.horizon_len
-    )
-    val_dataset = FinancialDataset(
-        series=val_series, config=config, context_length=128, horizon_length=config.horizon_len
-    )
-    print("Train samples:", len(train_dataset))
-    print("Val samples:", len(val_dataset))
-    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
-    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, collate_fn=collate_fn)
-
-    optimizer = optim.Adam(model.parameters(), lr=1e-4)
-
-    for epoch in range(num_epochs):
-        model.train()
-        total_train_loss = 0.0
-
-        for x_context, x_padding, freq, x_future in train_dataloader:
-            x_context, x_padding, freq, x_future = (
-                x_context.to(device),
-                x_padding.to(device),
-                freq.to(device),
-                x_future.to(device),
-            )
-            predictions = model(x_context, x_padding.float(), freq)
-            # predictions shape => [B, N, horizon_len, (1 + #quantiles)]
-            predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
-            last_patch_pred = predictions_mean[:, -1, :]  # => [B, horizon_len]
-
-            # x_future => [B, horizon_len]
-            loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
+            predictions = self.model(x_context, x_padding.float(), freq)
+            predictions_mean = predictions[..., 0]
+            last_patch_pred = predictions_mean[:, -1, :]
+            loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
 
             optimizer.zero_grad()
             loss.backward()
             optimizer.step()
 
-            total_train_loss += loss.item()
+            total_loss += loss.item()
 
-        avg_train_loss = total_train_loss / len(train_dataloader)
+        return total_loss / n_batches
 
-        # -------- Compute validation loss --------
-        model.eval()
-        total_val_loss = 0.0
-        with torch.no_grad():
-            for x_context, x_padding, freq, x_future in val_dataloader:
-                x_context, x_padding, freq, x_future = (
-                    x_context.to(device),
-                    x_padding.to(device),
-                    freq.to(device),
-                    x_future.to(device),
-                )
-                predictions = model(x_context, x_padding.float(), freq)
-                predictions_mean = predictions[..., 0]
-                last_patch_pred = predictions_mean[:, -1, :]
-                val_loss = torch.mean((last_patch_pred - x_future.squeeze(-1)) ** 2)
-                total_val_loss += val_loss.item()
+    @torch.no_grad()
+    def _validate(self, val_loader: DataLoader) -> float:
+        """Perform validation."""
+        self.model.eval()
+        total_loss = 0.0
 
-        avg_val_loss = total_val_loss / max(len(val_dataloader), 1)
+        for batch in val_loader:
+            x_context, x_padding, freq, x_future = [t.to(self.device) for t in batch]
 
-        print(f"[Epoch {epoch+1}] Train Loss: {avg_train_loss:.4f} | Val Loss: {avg_val_loss:.4f}")
+            predictions = self.model(x_context, x_padding.float(), freq)
+            predictions_mean = predictions[..., 0]
+            last_patch_pred = predictions_mean[:, -1, :]
 
-    torch.save(model.state_dict(), "timesfm_finetuned.pth")
-    return model, train_dataloader, val_dataloader
+            loss = self.loss_fn(last_patch_pred, x_future.squeeze(-1))
+            total_loss += loss.item()
 
+        return total_loss / len(val_loader)
 
-def plot_predictions(model, dataloader):
-    model.eval()
-    with torch.no_grad():
-        x_context, x_padding, freq, x_future = next(iter(dataloader))
-        x_context, x_padding, freq, x_future = (
-            x_context.to(device),
-            x_padding.to(device),
-            freq.to(device),
-            x_future.to(device),
-        )
-        # Forward pass
-        predictions = model(x_context, x_padding.float(), freq)
-        # => [B, N, horizon_len, (1 + #quantiles)]
-        predictions_mean = predictions[..., 0]  # => [B, N, horizon_len]
-        last_patch_prediction = predictions_mean[:, -1, :]  # => [B, horizon_len]
+    def finetune(self, train_dataset: Dataset, val_dataset: Dataset) -> Dict[str, Any]:
+        """
+        Finetune the TimesFM model on the provided datasets.
 
-        # We'll plot only the first sample in the batch
-        i = 0
-        pred_vals = last_patch_prediction[i].cpu().numpy()  # [horizon_len]
-        context_vals = x_context[i].cpu().numpy()  # [context_len]
-        future_vals = x_future[i].cpu().numpy()  # [horizon_len]
+        Args:
+            train_dataset: Training dataset
+            val_dataset: Validation dataset
 
-        horizon_len = future_vals.shape[0]
-        context_len = context_vals.shape[0]
+        Returns:
+            Dict containing training history and best model path
+        """
+        self.model = self.model.to(self.device)
 
-        plt.figure(figsize=(10, 5))
+        train_loader = self._create_dataloader(train_dataset, "train")
+        val_loader = self._create_dataloader(val_dataset, "val")
 
-        # Plot context
-        plt.plot(range(context_len), context_vals, label="Context (History)", color="blue")
-
-        # Plot predicted future
-        plt.plot(
-            range(context_len, context_len + horizon_len),
-            pred_vals,
-            label="Predicted Future",
-            color="orange",
+        optimizer = optim.Adam(
+            self.model.parameters(), lr=self.config.learning_rate, weight_decay=self.config.weight_decay
         )
 
-        # Plot ground truth future
-        plt.plot(
-            range(context_len, context_len + horizon_len),
-            future_vals,
-            label="Ground Truth Future",
-            color="green",
-            linestyle="--",
-        )
+        history = {"train_loss": [], "val_loss": [], "learning_rate": []}
 
-        plt.xlabel("Time")
-        plt.ylabel("Value")
-        plt.title("Model Forecast vs. Ground Truth")
-        plt.legend()
-        plt.show()
-        plt.savefig("pic_predictions.png")
+        self.logger.info(f"Starting training for {self.config.num_epochs} epochs...")
+        self.logger.info(f"Training samples: {len(train_dataset)}")
+        self.logger.info(f"Validation samples: {len(val_dataset)}")
 
+        try:
+            for epoch in range(self.config.num_epochs):
+                train_loss = self._train_epoch(train_loader, optimizer)
 
-if __name__ == "__main__":
-    # Example usage
-    model, train_dl, val_dl = train_model(
-        ticker="AAPL",
-        start="2012-01-01",
-        end="2019-01-01",
-        train_split=0.8,
-        batch_size=256,
-        num_epochs=50,
-        pretrained=True,
-    )
+                val_loss = self._validate(val_loader)
 
-    plot_predictions(model, val_dl)
+                current_lr = optimizer.param_groups[0]["lr"]
+
+                history["train_loss"].append(train_loss)
+                history["val_loss"].append(val_loss)
+                history["learning_rate"].append(current_lr)
+
+                metrics = {
+                    "train_loss": train_loss,
+                    "val_loss": val_loss,
+                    "learning_rate": current_lr,
+                    "epoch": epoch + 1,
+                }
+
+                if self.config.use_wandb:
+                    wandb.log(metrics)
+
+                print(f"[Epoch {epoch+1}] Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")
+        except KeyboardInterrupt:
+            self.logger.info("Training interrupted by user")
+
+        return {"history": history}