From 8bbf2194d21f01b5fa984be8c1e250406d804529 Mon Sep 17 00:00:00 2001 From: Diego Urgell Date: Tue, 2 Apr 2024 13:38:35 -0700 Subject: [PATCH] Fix compatibility with Python <= 3.8 (#172) Summary: Pull Request resolved: https://github.com/pytorch/torchsnapshot/pull/172 Reviewed By: JKSenthil Differential Revision: D55627601 fbshipit-source-id: a071b20ed363b617bd50d88620f78e76cb0a7214 --- torchsnapshot/io_preparers/sharded_tensor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/torchsnapshot/io_preparers/sharded_tensor.py b/torchsnapshot/io_preparers/sharded_tensor.py index b16d380..eb969c4 100644 --- a/torchsnapshot/io_preparers/sharded_tensor.py +++ b/torchsnapshot/io_preparers/sharded_tensor.py @@ -17,7 +17,7 @@ from dataclasses import dataclass from functools import reduce from operator import mul -from typing import Callable, List, Optional, Tuple +from typing import Callable, List, Optional, Tuple, Union import torch from torch.distributed._shard.sharded_tensor import ( @@ -199,7 +199,7 @@ def prepare_read( cls, entry: ShardedTensorEntry, obj_out: Optional[ShardedTensor] = None, - ) -> Tuple[List[ReadReq], Future[ShardedTensor | torch.Tensor]]: + ) -> Tuple[List[ReadReq], Future[Union[ShardedTensor, torch.Tensor]]]: # Note: in case obj_out is None, a Future[Tensor] will be returned if obj_out is None: obj_out = ShardedTensorIOPreparer.empty_tensor_from_sharded_tensor_entry(