Update README.md

README.md

@@ -33,11 +33,24 @@ You can load the model weights for inference or fine-tuning as follows:
 
 ```python
 import torch
-from pw_foucast import PWFouCast 
+from pw_foucast import PWFouCast
+from safetensors.torch import load_model
+from huggingface_hub import hf_hub_download
+
+MODEL_REGISTRY = {
+    'pw_foucast': PW_FouCast,
+}
+ModelClass = MODEL_REGISTRY.get(args.model.lower())
+
+model = ModelClass(**model_kwargs).to(args.device)
+model = torch.nn.DataParallel(model)
 
 # Load the model from Hugging Face
-model = PWFouCast.from_pretrained("your-username/pw-foucast")
+weights_path = hf_hub_download(repo_id=f"Onemiss/PW-FouCast/{args.model}/{args.dataset}", filename="model.safetensors")
+load_model(model, weights_path)
+
+# Eval
+model.eval()
+……
 
-state_dict = torch.hub.load_state_dict_from_url("[https://huggingface.co/your-username/pw-foucast/resolve/main/pytorch_model.bin](https://huggingface.co/your-username/pw-foucast/resolve/main/pytorch_model.bin)")
-model.load_state_dict(state_dict)
 ```