Typing fix

Author: fegin

torchft/checkpointing.py

@@ -172,7 +172,7 @@ def load_from_address(cls, address: str, timeout: timedelta) -> T:
             data = f.read()
 
         reader = io.BytesIO(data)
-        # We have to set weights_only to True as there are some non-tensor
+        # We have to set weights_only to False as there are some non-tensor
         # states like lr_scheduler.
         return torch.load(reader, weights_only=False)
 

torchft/manager.py

@@ -225,7 +225,7 @@ def __init__(
         self._participating_world_size: int = 0
 
     def set_state_dict_fns(
-        self, load_state_dict: Callable[T, None], state_dict: Callable[[], T]
+        self, load_state_dict: Callable[[T], None], state_dict: Callable[[], T]
     ) -> None:
         self._load_state_dict = load_state_dict
         self._user_state_dict = state_dict