diff --git a/docs/source-pytorch/extensions/callbacks_state.rst b/docs/source-pytorch/extensions/callbacks_state.rst
index 25f6ae4fb5..c83f3730fb 100644
--- a/docs/source-pytorch/extensions/callbacks_state.rst
+++ b/docs/source-pytorch/extensions/callbacks_state.rst
@@ -22,9 +22,9 @@ the following example.
             self.state = {"epochs": 0, "batches": 0}
 
         @property
-        def state_key(self):
+        def state_key(self) -> str:
             # note: we do not include `verbose` here on purpose
-            return self._generate_state_key(what=self.what)
+            return f"Counter[what={self.what}]"
 
         def on_train_epoch_end(self, *args, **kwargs):
             if self.what == "epochs":