@@ -1,4 +1,5 @@
 # Owner(s): ["oncall: distributed"]
+import contextlib
 import io
 import os
 import tempfile
@@ -98,6 +99,39 @@ def test_local_model(self):
         self.assertTrue(tracker._cur_module_name != "")
         self.assertTrue(hasattr(tracker, "_num_alloc_retries"))
 
+    def test_save_load_preserves_op_index(self):
+        """
+        A save/load round trip must restore the operation count so that
+        ``summary()`` keeps reporting the per-operator memory deltas. This runs
+        on CPU without an accelerator by populating the traces directly, mirror-
+        ing what ``_record_memory_stats()`` does during a monitored run.
+        """
+        tracker = MemoryTracker()
+        for i, name in enumerate(("op_a", "op_b")):
+            tracker.memories_allocated[i] = (name, float(i + 1))
+            tracker.memories_active[i] = (name, float(i + 1))
+            tracker.memories_reserved[i] = (name, float(i + 1))
+        tracker._op_index = 2
+
+        path = "memory_op_index.trace"
+        try:
+            tracker.save_stats(path)
+
+            loaded = MemoryTracker()
+            loaded.load(path)
+
+            self.assertEqual(loaded._op_index, tracker._op_index)
+
+            # With _op_index lost, summary() iterates range(1, 0) and emits no
+            # operator deltas; once restored, the operator name must appear.
+            buf = io.StringIO()
+            with contextlib.redirect_stdout(buf):
+                loaded.summary()
+            self.assertIn("op_b", buf.getvalue())
+        finally:
+            if os.path.exists(path):
+                os.remove(path)
+
 
 if __name__ == "__main__":
     run_tests()
@@ -212,6 +212,7 @@ def save_stats(self, path: str) -> None:
             "memories_reserved": self.memories_reserved,
             "markers": self._markers,
             "num_alloc_retries": self._num_alloc_retries,
+            "op_index": self._op_index,
         }
 
         with open(path, "wb") as f:
@@ -227,7 +228,10 @@ def load(self, path: str) -> None:
         self.memories_reserved = stats["memories_reserved"]
         self._markers = stats["markers"]
         self._num_alloc_retries = stats["num_alloc_retries"]
-        self._op_index = len(self.memories_allocated)
+        # Restore the operation count so summary()/show_traces() can bound their
+        # iteration. Stats files written before op_index was persisted do not
+        # carry it, so reconstruct it from the number of recorded operations.
+        self._op_index = stats.get("op_index", len(self.memories_allocated))
 
     def _create_pre_forward_hook(self, name: str) -> Callable:
         """Prefix operator name with current module and 'forward', and insert 'fw_start' marker at forward pass start."""