Update README.md

README.md

@@ -5,11 +5,11 @@ license: apache-2.0
 
 Code for [paper](https://arxiv.org/pdf/2403.03542) DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training (ICML'2024). It  pretrains neural operator transformers (from **7M** to **1B**)  on multiple PDE datasets. We will release the pre-trained weights soon.
 
-<!-- ![fig1](/resources/dpot.jpg) -->
+![fig1](./resources/dpot.jpg)
 
 Our pre-trained DPOT achieves the state-of-the-art performance on multiple PDE datasets and could be used for finetuning on different types of downstream PDE problems.
 
-<!-- ![fig2](/resources/dpot_result.jpg) -->
+![fig2](./resources/dpot_result.jpg)
 
 
 
@@ -26,8 +26,16 @@ We have five pre-trained checkpoints of different sizes.
 | Huge   | 2048          | 8092    | 27     | 8     | 1.03B      |
 
 
+#### Loading pre-trained model
 
-##### Dataset Protocol
+Here is an example code of loading pre-trained model.
+```
+model = DPOTNet(img_size=128, patch_size=8, mixing_type='afno', in_channels=4, in_timesteps=10, out_timesteps=1, out_channels=4, normalize=False, embed_dim=512, modes=32, depth=4, n_blocks=4, mlp_ratio=1, out_layer_dim=32, n_cls=12)
+model.load_state_dict(torch.load('model_Ti.pth')['model'])
+```
+
+
+#### Datasets
 
 All datasets are stored using hdf5 format, containing  `data`  field. Some datasets are stored with individual hdf5 files, others are stored within a single hdf5 file.