Update README.md

README.md

@@ -31,6 +31,40 @@ of the heat equation, given a pair of previous states.
 
 ---
 
+### How To Load Pretrained Models
+
+PDETransformer can be loaded via
+
+```python
+import torch
+from pdetransformer.core.mixed_channels import PDETransformer
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load fine-tuned model from Hugging Face
+model = PDETransformer.from_pretrained(
+    "saipuppala/pde-transformer-mc-s-diffusion-heat"
+).to(device)
+model.eval()
+
+# Example input: batch of size 1, 2 time steps, 64x64 grid
+# x[..., 0, :, :] ~ u(t0), x[..., 1, :, :] ~ u(t1)
+x = torch.randn((1, 2, 64, 64), dtype=torch.float32, device=device)
+
+with torch.no_grad():
+    out = model(x)                      # model output (tensor / dict / object)
+    pred_all = out if isinstance(out, torch.Tensor) else (
+        getattr(out, "prediction", None)
+        or getattr(out, "sample", None)
+        or next(v for v in out.values() if isinstance(v, torch.Tensor))
+    )
+
+# Convention: channel 1 corresponds to the next state prediction u(t2)
+u_t2_pred = pred_all[:, 1]             # shape: (B, H, W)
+print(u_t2_pred.shape)
+
+```
+
 ## 📝 Model Description
 
 PDE-Transformer is a transformer-based foundation model for physics simulations on regular grids.