Update README to include both shipped checkpoints

README.md

@@ -34,9 +34,10 @@ OmniMorph is a unified framework for 2D/3D multi-modal medical imaging (CT, MRI,
 | `Scripts/` | Auxiliary scripts (registration, evaluation) |
 | `tests/` | Pytest suite for `OMorpher` and loss functions |
 | `bash_*.sh`, `*.slurm` | SLURM submission scripts (CUDA + Intel XPU/Dawn) |
-| `Models/all_om_net/000110_all_om_net.pth` | Trained checkpoint (epoch 110, multi-modal `recmulmodmutattnnet`) |
+| `Models/all_om_net/000110_all_om_net.pth` | Trained checkpoint — production multi-modal `recmulmodmutattnnet` (epoch 110, ~3.0 GB) |
+| `Models/all_recmulmodmutattnnet/000010_all_recmulmodmutattnnet.pth` | Earlier `recmulmodmutattnnet` run (epoch 10, ~906 MB) |
 
-> **Note** Only the final checkpoint (epoch 110) is shipped here. Earlier epochs and the `bert_large_uncased` weights are not bundled — download `bert-large-uncased` from the official Hugging Face repo if you need the contrastive text encoder.
+> **Note** Only the final checkpoint of each training run is shipped — intermediate epochs and the `bert_large_uncased` weights are not bundled. Download `bert-large-uncased` from the official Hugging Face repo if you need the contrastive text encoder.
 
 ## Setup
 
@@ -85,9 +86,14 @@ python OM_reg_flexres.py -C Config/config_om.yaml
 import torch
 from Diffusion.networks import get_net
 
-# Production network (multi-modal recmutattnnet)
+# Production network (multi-modal recmulmodmutattnnet)
 net = get_net("recmulmodmutattnnet")
-state = torch.load("Models/all_om_net/000110_all_om_net.pth", map_location="cpu")
+
+# Production checkpoint (epoch 110)
+ckpt_path = "Models/all_om_net/000110_all_om_net.pth"
+# Or earlier run: "Models/all_recmulmodmutattnnet/000010_all_recmulmodmutattnnet.pth"
+
+state = torch.load(ckpt_path, map_location="cpu")
 net.load_state_dict(state["model"] if "model" in state else state)
 net.eval()
 ```