wehub-resource-sync
27 lines
639 B
Python
27 lines
639 B
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
|
|
|
|
def check_and_handle_empty_buffer(
|
|
buffer_m: torch.Tensor,
|
|
original_shape: torch.Size,
|
|
original_size: int,
|
|
worker_error: torch.Tensor,
|
|
server_error: torch.Tensor,
|
|
) -> Optional[torch.Tensor]:
|
|
if original_size == 0:
|
|
if worker_error.numel():
|
|
worker_error.zero_()
|
|
if server_error.numel():
|
|
server_error.zero_()
|
|
if len(original_shape) > 1:
|
|
return buffer_m.reshape(original_shape)
|
|
return buffer_m
|
|
return None
|