github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-09 05:13:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

'fix/format' (#573)

Browse Source
This commit is contained in:
yuxuan-lou 2022-03-31 15:46:11 +08:00 committed by binmakeswell
parent b0f708dfc1
commit cfb41297ff
2 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
colossalai/engine/schedule/_base_schedule.py
View File

@ -85,8 +85,7 @@ class BaseSchedule(ABC):
data_iter: Iterable, data_iter: Iterable,
forward_only: bool, forward_only: bool,
return_loss: bool = True, return_loss: bool = True,
return_output_label: bool = True return_output_label: bool = True):
):
"""The process function over a batch of dataset for training or evaluation. """The process function over a batch of dataset for training or evaluation.
Args: Args:
@ -107,8 +106,9 @@ class BaseSchedule(ABC):
@staticmethod @staticmethod
def _call_engine_criterion(engine, outputs, labels): def _call_engine_criterion(engine, outputs, labels):
assert isinstance(outputs, (torch.Tensor, list, tuple) assert isinstance(
), f'Expect output of model is (torch.Tensor, list, tuple), got {type(outputs)}' outputs,
(torch.Tensor, list, tuple)), f'Expect output of model is (torch.Tensor, list, tuple), got {type(outputs)}'
if isinstance(outputs, torch.Tensor): if isinstance(outputs, torch.Tensor):
outputs = (outputs,) outputs = (outputs,)
if isinstance(labels, torch.Tensor): if isinstance(labels, torch.Tensor):