Update streaming/base/storage/download.py

mosaicml · Dec 9, 2024 · dce3d6e · dce3d6e
1 parent 3e64719
commit dce3d6e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/streaming/base/storage/download.py b/streaming/base/storage/download.py
@@ -266,7 +266,7 @@ def __getstate__(self) -> dict:
         state['_s3_client'] = None  # Exclude _s3_client from being pickled
         return state
 
-    def __setstate__(self, state):
+    def __setstate__(self, state: dict):
         self.__dict__.update(state)
         self._s3_client = None  # Ensure _s3_client is reset after unpickling