update warning

pytorch · Apr 16, 2024 · 40c6b93 · 40c6b93
1 parent 92071f1
commit 40c6b93
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git 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
@@ -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."
                         )