add typing to utilities/memory.py

Author: yicongd

pyproject.toml

@@ -96,7 +96,6 @@ module = [
     "pytorch_lightning.utilities.data",
     "pytorch_lightning.utilities.distributed",
     "pytorch_lightning.utilities.fetching",
-    "pytorch_lightning.utilities.memory",
     "pytorch_lightning.utilities.meta",
 ]
 ignore_errors = "True"

pytorch_lightning/utilities/memory.py

@@ -17,7 +17,7 @@
 import os
 import shutil
 import subprocess
-from typing import Any, Dict
+from typing import Any, cast, Dict, IO
 
 import torch
 from torch.nn import Module
@@ -184,6 +184,6 @@ def get_model_size_mb(model: Module) -> float:
         Number of megabytes in the parameters of the input module.
     """
     model_size = _ByteCounter()
-    torch.save(model.state_dict(), model_size)
+    torch.save(model.state_dict(), cast(IO[bytes], model_size))
     size_mb = model_size.nbytes / 1e6
     return size_mb