From 97c1aab05873f1027467e535f9d2aecbe2a748a0 Mon Sep 17 00:00:00 2001 From: wandz Date: Wed, 7 Aug 2024 19:18:52 +0800 Subject: [PATCH] Fix the incorrect statement --- megatron/core/transformer/transformer_block.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/transformer/transformer_block.py b/megatron/core/transformer/transformer_block.py index fbcb2d72c1..0b0274906a 100755 --- a/megatron/core/transformer/transformer_block.py +++ b/megatron/core/transformer/transformer_block.py @@ -397,7 +397,7 @@ def forward( else: fp8_context = nullcontext() - with rng_context and fp8_context: + with rng_context, fp8_context: # Forward pass. if self.config.recompute_granularity == 'full' and self.training: hidden_states = self._checkpointed_forward(