From 318fbf1145d358ac3b4405ceb56954a94116d089 Mon Sep 17 00:00:00 2001 From: Kirigaya Kazuto <59416203+LSTM-Kirigaya@users.noreply.github.com> Date: Thu, 8 Sep 2022 22:04:34 +0800 Subject: [PATCH] [NFC] polish colossalai/utils/multi_tensor_apply/multi_tensor_apply.py code style (#1559) --- colossalai/pipeline/rpc/PipelineBase.py | 2 +- colossalai/utils/multi_tensor_apply/multi_tensor_apply.py | 8 ++------ 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/colossalai/pipeline/rpc/PipelineBase.py b/colossalai/pipeline/rpc/PipelineBase.py index 34c428b9d..8cc90b825 100644 --- a/colossalai/pipeline/rpc/PipelineBase.py +++ b/colossalai/pipeline/rpc/PipelineBase.py @@ -778,4 +778,4 @@ class OneFOneBPipelineEngine(PipelineEngineBase): criterion: Callable = None, checkpoint: bool = False) -> None: use_1F1B = True - super().__init__(module_partitions, stage_num, num_microbatches, device, use_1F1B, chunk, criterion, checkpoint) \ No newline at end of file + super().__init__(module_partitions, stage_num, num_microbatches, device, use_1F1B, chunk, criterion, checkpoint) diff --git a/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py b/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py index 1de82a414..4e847f17b 100644 --- a/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py +++ b/colossalai/utils/multi_tensor_apply/multi_tensor_apply.py @@ -26,13 +26,9 @@ class MultiTensorApply(object): raise RuntimeError( "Attempted to call MultiTensorApply method, but MultiTensorApply " "is not available, possibly because Apex was installed without " - "--cpp_ext --cuda_ext. Original import error message:", - MultiTensorApply.import_err) + "--cpp_ext --cuda_ext. Original import error message:", MultiTensorApply.import_err) def __call__(self, op, noop_flag_buffer, tensor_lists, *args): self.check_avail() - return op(self.chunk_size, - noop_flag_buffer, - tensor_lists, - *args) + return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args)