Update README.md

README.md

@@ -31,7 +31,7 @@ of the heat equation, given a pair of previous states.
 
 ---
 
-### How To Load Pretrained Models
+### How To Load Model
 
 PDETransformer can be loaded via
 
@@ -63,6 +63,25 @@ with torch.no_grad():
 u_t2_pred = pred_all[:, 1]             # shape: (B, H, W)
 print(u_t2_pred.shape)
 
+```
+
+### How To Download the Model
+
+PDETransformer can be loaded via
+
+```python
+from huggingface_hub import hf_hub_download
+import torch
+
+ckpt_path = hf_hub_download(
+    repo_id="saipuppala/diffusion_transformer ",
+    filename="diffusion_finetuned.pth",
+)
+
+state_dict = torch.load(ckpt_path, map_location="cpu")
+# then load into a PDETransformer instance as shown above
+
+
 ```
 
 ## 📝 Model Description