Update docstring for model loading code

Author: animan42

src/cube/inference/utils.py

@@ -37,21 +37,20 @@ def parse_structured(cfg_type: Any, cfg: DictConfig) -> Any:
     return scfg
 
 
-def load_model_weights(model: torch.nn.Module, ckpt_path: str) -> torch.nn.Module:
+def load_model_weights(model: torch.nn.Module, ckpt_path: str) -> None:
     """
     Load a safetensors checkpoint into a PyTorch model.
+    The model is updated in place.
 
     Args:
         model: PyTorch model to load weights into
         ckpt_path: Path to the safetensors checkpoint file
 
     Returns:
-        The model with loaded weights
+        None
     """
-    assert ckpt_path.endswith(".safetensors"), (
-        f"Checkpoint path '{ckpt_path}' is not a safetensors file"
-    )
+    assert ckpt_path.endswith(
+        ".safetensors"
+    ), f"Checkpoint path '{ckpt_path}' is not a safetensors file"
 
     load_model(model, ckpt_path)
-
-    return model