[Fix] Type annotation correction for UpdateWeightsFromTensorReqInput (#4532)
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
This commit is contained in:
@@ -540,7 +540,8 @@ class UpdateWeightsFromDistributedReqOutput:
|
||||
|
||||
@dataclass
|
||||
class UpdateWeightsFromTensorReqInput:
|
||||
serialized_named_tensors: bytes # indeed Dict[str, torch.Tensor]
|
||||
# List containing one serialized Dict[str, torch.Tensor] per TP worker
|
||||
serialized_named_tensors: List[bytes]
|
||||
load_format: Optional[str]
|
||||
flush_cache: bool
|
||||
|
||||
|
||||
Reference in New Issue
Block a user