From 40c6b9339f74b5a5094bf8de58d39d3b40fdfd4b Mon Sep 17 00:00:00 2001 From: "Liao, Wei" Date: Tue, 16 Apr 2024 11:20:07 +0800 Subject: [PATCH] update warning --- torchbenchmark/util/extra_args.py | 2 +- torchbenchmark/util/model.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/torchbenchmark/util/extra_args.py b/torchbenchmark/util/extra_args.py index c9f2b693ef..3c615fbc88 100644 --- a/torchbenchmark/util/extra_args.py +++ b/torchbenchmark/util/extra_args.py @@ -189,7 +189,7 @@ def apply_decoration_args( elif model.test == "train": # the model must implement staged train test warnings.warn( - "Usually models do not want to enable amp only in forward path, so expected " + "Usually models only want to enable AMP in forward path, so expected " "model to have staged train support." ) import torch diff --git a/torchbenchmark/util/model.py b/torchbenchmark/util/model.py index 3a1d76a84e..a9a1effbdf 100644 --- a/torchbenchmark/util/model.py +++ b/torchbenchmark/util/model.py @@ -490,7 +490,7 @@ def enable_amp(self): self.add_context(self.amp_context, TEST_STAGE.FORWARD) else: warnings.warn( - "Usually models do not want to enable amp only in forward path, so expected " + "Usually models only want to enable AMP in forward path, so expected " "model to have staged train support." )