@@ -4201,7 +4201,7 @@ def _prepare_for_write_opaque(
4201
4201
w .set_min_file_version_number (OPAQUE_TENSORIZER_VERSION )
4202
4202
4203
4203
@staticmethod
4204
- def _do_clone (write_spec , dependency : Optional [concurrent . futures . Future ]):
4204
+ def _do_clone (write_spec , dependency : Optional [_Future ]):
4205
4205
if dependency is not None :
4206
4206
dependency .result (_TIMEOUT )
4207
4207
write_spec .tensor = write_spec .tensor .clone ().detach ()
@@ -4234,7 +4234,7 @@ def _prepare_for_write_encryption(
4234
4234
]
4235
4235
)
4236
4236
4237
- clone_tasks = []
4237
+ clone_tasks : List [ _Future ] = []
4238
4238
for w in shared_write_specs [1 :]:
4239
4239
clone_tasks .append (
4240
4240
self ._computation_pool .submit (
@@ -4401,7 +4401,7 @@ def compute_sha256(
4401
4401
for w in write_specs :
4402
4402
old_tensor_data_task = w .tensor_data_task
4403
4403
4404
- hash_tasks = []
4404
+ hash_tasks : List [ _Future ] = []
4405
4405
if w .include_crc32 :
4406
4406
crc32_task = self ._computation_pool .submit (
4407
4407
compute_crc32 , w , old_tensor_data_task
@@ -4418,7 +4418,7 @@ def compute_sha256(
4418
4418
self ._jobs .extend (hash_tasks )
4419
4419
4420
4420
def _do_encryption (self , write_specs : Sequence [_WriteSpec ]) -> None :
4421
- def encrypt (write_spec , dependency : _Future ):
4421
+ def encrypt (write_spec , dependency : Optional [ _Future ] ):
4422
4422
if dependency is not None :
4423
4423
dependency .result (_TIMEOUT )
4424
4424
try :
0 commit comments